Docker로 Tensorflow GPU 실행하기

NVIDIA 그래픽 드라이버 설치

---

우선 아래 페이지에 들어가 자신에게 맞는 NVIDIA 그래픽 드라이버를 받고 설치해 준다.

https://www.nvidia.co.kr/Download/index.aspx?lang=kr

나는 가끔 게임을 하므로, Game Ready 드라이버(GRD)를 다운받아 설치했다.

설치가 정상적으로 완료되었다면 GPU 정보를 확인하기 위해 nvidia-smi 명령어를 사용할 수 있다.

nvidia-smi

NVIDIA Docker 설치

---

먼저 아래 글에서 설명한 내용을 토대로 Docker를 설치해준다.
https://sirius-mhlee.github.io/posts/wsl-docker-install/

1. NVIDIA Docker를 설치하기 전에 Docker 서비스를 종료해준다.

sudo service docker stop

2. NVIDIA의 Official GPG Key를 등록하고 Repository도 APT에 등록한다.

curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
  && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
  sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
  sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
sudo apt update

3. NVIDIA Container Toolkit을 설치한다.

sudo apt install nvidia-container-toolkit

4. 종료했던 Docker 서비스를 다시 시작한다.

sudo service docker start

5. Docker 컨테이너에서 GPU 정보를 확인한다.

docker run --rm --gpus all ubuntu nvidia-smi

컨테이너 실행 시 --gpus 옵션을 사용하여 GPU를 사용할 수 있도록 한다.

• --gpus all : 컨테이너에서 모든 GPU 사용

Docker 컨테이너에서 GPU를 정상적으로 인식하고 있는 것을 확인할 수 있다.

Tensorflow GPU 컨테이너 실행

---

Docker Hub의 Tensorflow 공식 이미지를 이용하여 컨테이너를 시작하게 된다.

docker run --name tf_gpu_jupyter --gpus all -d -p 8888:8888 tensorflow/tensorflow:latest-gpu-jupyter

Docker의 logs 명령으로 Jupyter Notebook의 접속 주소를 확인한다.

docker logs <container>

• docker logs <container> : 지정된 컨테이너에서 출력한 내용을 확인

이제 Jupyter Notebook으로 접속해서 Tensorflow가 GPU를 제대로 인식하고 있는지 확인해보자.

from tensorflow.python.client import device_lib
device_lib.list_local_devices()

Tensorflow에서 GPU를 제대로 인식하는 것을 확인할 수 있다.

Jupyter Notebook이 아닌 스크립트 버전을 실행하려면 다음 명령어를 사용하면 된다.

docker run --name tf_gpu --gpus all -i -t -v ~/tf_gpu_data:/root tensorflow/tensorflow:latest-gpu