[Docker] Tensorflow GPU 사용하기

ubuntu 이미지를 깔아서 그 위에 텐서플로우를 설치하여도 되지만,

텐서플로우 공식 홈페이지에 쥬피터와 텐서플로우가 모두 깔려져 있는 이미지가 이미 존재하길래 해당 이미지를 사용해주었습니다.

 

출처 : https://www.tensorflow.org/install/docker?hl=ko#download_a_tensorflow_docker_image 

Docker  |  TensorFlow

 

 

 

그냥 latest 태그를 깔게 되면 CPU만을 사용하는 텐서플로우를 깔게 됩니다.

저는 쥬피터도 사용하고 GPU도 사용할꺼라서 latest-gpu-jupyter를 깔게 되었습니다.

 

공식 홈페이지에 있는 코드를 복사하여 사용할 수도 있지만 아래와 같이 Docker Hub에 들어가서 어떤 태그가 있는지 살펴보고,

본인에게 맞는 태그를 설치하면 됩니다.

 

 

보니깐 쥬피터 공식 페이지에서도 텐서플로우 docker 이미지를 제공하네요.

 

 

 

sudo docker pull tensorflow/tensorflow:latest-gpu-jupyter

 

이미지를 사용하여 컨테이너를 실행시켜 줍니다.

 

sudo docker run \
    --rm \
    -it \
    --name boostcourse \
    --gpus all \
    -v /home/scone/deepdish/boostcourse/:/tf/boostcourse \
    -p 8888:8888 \
    tensorflow/tensorflow:latest-gpu-jupyter

차례대로 설명 드리면

 

• --rm
  • 컨테이너 종료시 자동 컨테이너 삭제
    백그라운드에 컨테이너가 남아서 리소스를 잡아먹지 않겠끔 바로바로 삭제해주는 옵션입니다.
    매우 편한데 문제는 commit 안하고 종료시 작업했던 데이터가 날라갈 수 있습니다.
• -it
  • 코드 작성 이후 terminal를 종료하지 않고, 로그를 계속 띄울 수 있겠끔 해주는 옵션 입니다.
  • 보통 bash 창이 뜨겠끔 설정해놓는데 여기서는 바로 쥬피터가 실행되므로 필요시 attech 옵션을 통해 terminal를 별도로 띄워주어야 합니다.
• --name boostcourse
  • container 이름을 지정합니다.
• --gpus all
  • gpu를 모두 다 사용하겠다는 옵션 입니다.
• -v {local 파일 경로}:{container 내의 파일 시스템 경로}
  • local과 공유할 volume을 설정해줍니다.
  • 여기다가 파일을 쥬피터 파일을 작업하고, 저장해두기 때문에 --rm 옵션을 자유롭게 써도 문제가 되지 않는 것 같습니다.
  • 라이브러리를 저장해두고 싶다면 commit을 반드시 하도록 하고, 그냥 매번 리셋하고 싶다면 쥬피터 파일만 잘 저장해두고 터미널을 그냥 끄면 알아서 컨테이너가 삭제돼서 편합니다.
• -p 8888:8888
  • 8888 로컬 포트와 8888 컨테이너 포트를 연결해주었습니다.
  • 초기에는 root 사용자가 접속 가능하겠끔 설정해주어야합니다.
    해당 내용은 다음 게시글에서 확인할 수 있습니다. https://callmescone.tistory.com/332
• 위에서 pull로 받아온 이미지 이름 작성

 

 

TensorFlow가 GPU를 잘 인식하는지 확인

• 디바이스 리스트를 먼저 봐줍니다.

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))

CPU 하나와 GPU 하나를 잘 인식하고 있습니다.

 

• GPU를 사용하겠끔 하는 코드 입니다.

import os

os.environ["CUDA_VISIBLE_DEVICES"]="0"
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        tf.config.experimental.set_memory_growth(gpus[0], True)
    except RuntimeError as e:
        print(e)

빈 리스트가 나온다면 장치에 gpu가 없다는 의미 입니다.

 

• CUDA 체크

import tensorflow as tf
print(tf.test.is_built_with_cuda())

CUDA 또한 준비되어 있습니다.

 

 

watch -d -n 0.5 nvidia-smi

간단한 CNN 모델을 돌렸고 GPU가 사용되는것을 확인하였습니다.

https://callmescone.tistory.com/309

리눅스 시스템 모니터링 feat GPU 상태 확인 (watch)