nvidia-docker2 환경 기반 tensorflow 설치

이전 글에서 nvidia-docker2 설정 관련 내용들을 다루었습니다. 이제 docker hub에서 필요한 이미지를 받아서, tensorflow를 마저 세팅해보겠습니다.

2022.03.24 - [Machine learning] - NVIDIA container toolkit 설치

2022.03.23 - [Machine learning] - [nvidia driver 에러] error running hook: exit status 1, stdout: , stderr: nvidia-container-cli: initialization error: nvml error: driver not loaded: unknown.

 

tensorflow를 설치하기 전에, 먼저 cuda 버전과 cuDNN버전 matching을 먼저 확인해보았습니다. 종종 tensorflow를 설치하다가, version mismatch로 인해 다시 설치한 적이 몇번있기에 이번에는 설치전에 버전 관련 내용을 먼저 확인해보았습니다. 해당 내용은 tensorflow 홈페이지에 tested build configurations란에 나와있습니다. 저는 GPU를 사용할 것이기에, 관련 버전 matching을 확인해보면 tensorflow 2.8.0 기준으로는 cuDNN이 8.1버전이고 CUDA가 11.2까지 지원됩니다.

tensorflow GPU버전 matching (source: tensorflow 홈페이지)

https://www.tensorflow.org/install/source#gpu

소스에서 빌드  |  TensorFlow

 

docker hub에서 tensorflow 관련 image들을 찾아봅니다. 이번에 저는 tensorflow에서 제공하는 tensorflow2.8.0-gpu-jupyter를 받았습니다.

tensorflow 2.8.0 gpu jupyter image

해당 이미지에 대한 layer들을 살펴보면, cuda 세팅이나 cuDNN세팅이 모두 잘 되어 있는 것을 확인할 수 있습니다. 아래의 명령어로 해당 이미지를 받습니다.

docker pull tensorflow/tensorflow:2.8.0-gpu-jupyter

그러면 아래와 같이 해당 도커 이미지를 받습니다.

docker image pulling

이미지를 받은 후 sudo docker images를 커맨드 창에 입력해주면, 아래와 같이 다운받은 이미지를 확인할 수 있습니다.

sudo docker images

tensorflow 2.8.0-gpu-jupyter doker image 확인

docker가 있으니, 환경 설정에 들이는 시간이 엄청 단축되는 느낌입니다. 해당 이미지가 잘 잘동하는지 테스트해보겠습니다.

docker run --gpus all -it --rm tensorflow/tensorflow:2.8.0-gpu-jupyter \
python -c "import tensorflow as tf; print(tf.reduce_sum(tf.random.normal([1000, 1000])))"

아래와 같이 로그가 찍힌 것을 확인하였습니다. 받은 이미지가 정상적으로 작동하는 것을 확인할 수 있었습니다.

tensorflow docker image test