Visual Studio Codeのdevcontainerを利用すれば，自分好みの開発環境をDockerfileにまとめられるので便利です．今回は，PythonでTensorFlowを高速に実行するために，devcontainerでGPUを利用できるようにします．
始める前に，DockerのホストにCUDAのドライバーなど必要な準備をしておきます．
DockerでCUDAを使用するための準備 - みーのぺーじ
Dockerfileの作成
ベースイメージの選定
TensorFlow 2.9では，CPUを利用する時はoneDNN，GPUを利用する時はcuDNNを使用します*1．cuDNNは，libcudnn8というパッケージがubuntu 20.04用に配布されており，aptでインストールして使用します*2．
CUDAのインストール
NVIDIAが提供しているライブラリは以下に公開されていますので，適宜aptに登録して使用します*3．
cuDNNのインストール
2022/06/21現在利用可能なバージョンは，CUDA_VERSION 11.6.2, CUDNN_VERSION = 8.4.0.27 のようなので，パッケージ名はlibcudnn8=8.4.0.27-1+cuda11.6となります．
ubuntuをdevcontainerとして使用する
Visual Studio Codeの関連として提供されているlibrary-scripts/common-debian.shを実行すればよいみたいです*4．
これらを合体させて，以下のようなDockerfileを作成しました．
FROM ubuntu:20.04 RUN apt-get update && export DEBIAN_FRONTEND=noninteractive \ && apt-get -y install --no-install-recommends \ gnupg2 curl ca-certificates # Options for setup script ARG INSTALL_ZSH="true" ARG UPGRADE_PACKAGES="true" ARG USERNAME=vscode ARG USER_UID=1000 ARG USER_GID=$USER_UID # COPY .devcontainer/library-scripts/*.sh .devcontainer/library-scripts/*.env /tmp/library-scripts/ RUN yes | unminimize 2>&1 \ && bash /tmp/library-scripts/common-debian.sh "${INSTALL_ZSH}" "${USERNAME}" "${USER_UID}" "${USER_GID}" "${UPGRADE_PACKAGES}" "true" "true" \ && apt-get clean -y && rm -rf /var/lib/apt/lists/* /tmp/library-scripts # RUN curl -fsSLO https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb && \ dpkg -i cuda-keyring_1.0-1_all.deb && \ apt-get purge --autoremove -y curl \ && rm -rf /var/lib/apt/lists/* ENV NVIDIA_REQUIRE_CUDA "cuda>=11.6 brand=tesla,driver>=418,driver<419 brand=tesla,driver>=450,driver<451 brand=tesla,driver>=470,driver<471 brand=unknown,driver>=470,driver<471 brand=nvidia,driver>=470,driver<471 brand=nvidiartx,driver>=470,driver<471 brand=quadrortx,driver>=470,driver<471" ENV NV_CUDA_CUDART_VERSION 11.6.55-1 ENV NV_CUDA_COMPAT_PACKAGE cuda-compat-11-6 ENV CUDA_VERSION 11.6.2 ENV NV_CUDA_LIB_VERSION 11.6.2-1 ENV NV_NVTX_VERSION 11.6.124-1 ENV NV_LIBNPP_VERSION 11.6.3.124-1 ENV NV_LIBNPP_PACKAGE libnpp-11-6=${NV_LIBNPP_VERSION} ENV NV_LIBCUSPARSE_VERSION 11.7.2.124-1 ENV NV_LIBCUBLAS_PACKAGE_NAME libcublas-11-6 ENV NV_LIBCUBLAS_VERSION 11.9.2.110-1 ENV NV_LIBCUBLAS_PACKAGE ${NV_LIBCUBLAS_PACKAGE_NAME}=${NV_LIBCUBLAS_VERSION} ENV NV_LIBNCCL_PACKAGE_NAME libnccl2 ENV NV_LIBNCCL_PACKAGE_VERSION 2.12.10-1 ENV NCCL_VERSION 2.12.10-1 ENV NV_LIBNCCL_PACKAGE ${NV_LIBNCCL_PACKAGE_NAME}=${NV_LIBNCCL_PACKAGE_VERSION}+cuda11.6 ENV NV_CUDNN_VERSION 8.4.0.27 ENV NV_CUDNN_PACKAGE_NAME "libcudnn8" ENV NV_CUDNN_PACKAGE "libcudnn8=$NV_CUDNN_VERSION-1+cuda11.6" RUN apt-get update && apt-get install -y --no-install-recommends \ cuda-cudart-11-6=${NV_CUDA_CUDART_VERSION} \ ${NV_CUDA_COMPAT_PACKAGE} \ cuda-libraries-11-6=${NV_CUDA_LIB_VERSION} \ ${NV_LIBNPP_PACKAGE} \ cuda-nvtx-11-6=${NV_NVTX_VERSION} \ libcusparse-11-6=${NV_LIBCUSPARSE_VERSION} \ ${NV_LIBCUBLAS_PACKAGE} \ ${NV_LIBNCCL_PACKAGE} \ ${NV_CUDNN_PACKAGE} \ && apt-mark hold ${NV_CUDNN_PACKAGE_NAME} \ && apt-mark hold ${NV_LIBCUBLAS_PACKAGE_NAME} ${NV_LIBNCCL_PACKAGE_NAME} \ && ln -s cuda-11.6 /usr/local/cuda && \ rm -rf /var/lib/apt/lists/* RUN apt-get update && apt-get install -y --no-install-recommends \ python3 python3-pip && \ rm -rf /var/lib/apt/lists/* RUN pip3 --disable-pip-version-check --no-cache-dir install flake8 mypy black COPY requirements.txt /tmp/pip-tmp/ RUN pip3 --disable-pip-version-check --no-cache-dir install -r /tmp/pip-tmp/requirements.txt \ && rm -rf /tmp/pip-tmp
また，CUDAをdockerから使用するために，コンテナ実行時に--gpus=allオプションを指定します．これはdevcontainer.jsonに以下を追加します．
{ "runArgs": [ "--gpus=all" ] }
requirements.txtは以下のようにしました．
numpy==1.22.4 tensorflow==2.9.1 protobuf==3.19.4
動作確認
devcontainerに接続したターミナルでnvidia-smiを実行し，GPUの情報を取得してみます．
$ nvidia-smi
Tue Jun 21 13:35:59 2022
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 515.48.07 Driver Version: 515.48.07 CUDA Version: 11.7 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 NVIDIA GeForce ... On | 00000000:01:00.0 Off | N/A |
| N/A 31C P8 N/A / 30W | 1MiB / 2048MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
| No running processes found |
+-----------------------------------------------------------------------------+
TensorFlowからGPUの情報を取得してみます．
$ python3 Python 3.8.10 (default, Mar 15 2022, 12:22:08) [GCC 9.4.0] on linux Type "help", "copyright", "credits" or "license" for more information. >>> import tensorflow 2022-06-21 13:50:15.359370: I tensorflow/core/util/util.cc:169] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. >>> gpus = tensorflow.config.list_physical_devices(device_type = 'GPU') >>> len(gpus) 1 >>> gpus[0].name '/physical_device:GPU:0'
devcontainerから，cuDNNを利用したTensorFlowが使用できるようになりました．
*1:What's new in TensorFlow 2.9? — The TensorFlow Blog
*2:Installation Guide :: NVIDIA Deep Learning cuDNN Documentation
*3:dist/11.6.2/ubuntu2004/base/Dockerfile · master · nvidia / container-images / cuda · GitLab, dist/11.7.0/ubuntu2204/base/Dockerfile · master · nvidia / container-images / cuda · GitLab
*4:vscode-dev-containers/base.Dockerfile at v0.231.6 · microsoft/vscode-dev-containers · GitHub
