はじめに

CUDA on WSL2が発表され、また在宅勤務でリモートワークが必須になってきたので、これを機にコンテナを利用したリモート開発環境を整えようっていうことで、関連記事に記載したような形で環境を整えてきました。そこで本記事では、CUDA on WSL2上のCUDA コンテナに対しRemote containerの設定を行いVSCodeで接続・TensorFlowを実行することころまで実施します。
またAnacondaの規約が変更され有償になり、その次善策を検討する必要があったので、あわせてTensorFlowの過去バージョンが、どの程度利用できるかについても調べてみました。

結果(サマリ)

• 公開されているコンテナイメージを、テンプレートに当てはめるだけで環境構築が可能。
• コンテナ内のCUDAバージョンには依存せず、結構昔(TensorFlow v1.5, CUDA v9.0)のも動作する(≒バージョン切り替えが簡単)。
• ただしメモリへのアクセスに問題があるようで、学習サイズが小さくなればなるほど性能が低下していく。

検証環境

検証環境は下記になっています。ソフトウェアは2020/12/22時点で可能な限り最新にしています。

• ハードウェア
  • Intel Core i7 9700
  • DDR4-2666 16 GB (Dual Channel)
  • Geforce RTX 2070 SUPER (8GB)
• ソフトウェア
  • Windows 10 Pro Insider Preview, Build. 21277.1000
  • NVIDIA Driver for Windows 10: 465.21
  • WSL2 Kernel version 5.4.72
  • Docker version 20.10.1

環境構築

CUDA on WSL2の環境構築については、公式ページ、または下記の佐々木さんのQiita記事を参考にしてください。

• 待ってました CUDA on WSL 2 - Qiita

またリモート開発環境の構築は、下記の以前の記事を参考にしてください。

• VSCode+WSL2+Dockerによるリモート開発環境の構築) - Qiita

環境構築の概要をざっくり記載すると以下の通りです。

1. Windows Insider Programに参加
2. 「Windows → 設定 → 更新とセキュリティ → Windows Insider Program」より「Dev チャネル」に設定
3. Geforce Driver (Preview)をダウンロードし、インストール
4. WSL2を本手順に従いインストール、Microsoft StoreよりUbuntu18.04を導入
5. WSL2上のUbuntuにCUDA Toolkit、及びDocker、NVIDIA Container Toolkitをインストール(下記コマンドを参照)
6. リモート開発環境の設定(上記の別記事)

Ubuntu(WSL2)

# CUDA Toolkitのレポジトリを追加
$ sudo apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
$ sudo sh -c 'echo "deb http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64 /" > /etc/apt/sources.list.d/cuda.list'
$ sudo apt-get update

# CUDA Toolkitをインストール
$ sudo apt-get install -y cuda-toolkit-11-0

# Docker CEをインストール
$ sudo -E curl https://get.docker.com | sh

# NVIDIA Container Toolkitのレポジトリを追加
$ distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
$ curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
$ curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
$ curl -s -L https://nvidia.github.io/libnvidia-container/experimental/$distribution/libnvidia-container-experimental.list | sudo tee /etc/apt/sources.list.d/libnvidia-container-experimental.list

# NVIDIA Container Toolkitをインストール
$ sudo apt-get update
$ sudo apt-get install -y nvidia-docker2

# Dockerの再起動
$ sudo service docker stop
$ sudo service docker start

リモート開発用コンテナの準備

独自のリモート開発用コンテナの作成は、githubにあるテンプレートを参考にしながら、Dockerfileとdevcontainer.jsonを設定します。
作成したDockerfileとdevcontainer.jsonは下記の通りです。Dockerfile中のcommon-debian.shは.devcontainer/library-scriptsにテンプレートからコピーしてください。
コンテナのイメージは、Docker HubのTensorFlowを利用しています。

Dockerfile

ARG VARIANT="latest"
FROM tensorflow/tensorflow:${VARIANT}-gpu-jupyter

ARG INSTALL_ZSH="true"
ARG UPGRADE_PACKAGES="false"

ARG USERNAME=vscode
ARG USER_UID=1000
ARG USER_GID=$USER_UID
COPY library-scripts/*.sh /tmp/library-scripts/
RUN apt-get update && export DEBIAN_FRONTEND=noninteractive \
    && /bin/bash /tmp/library-scripts/common-debian.sh "${INSTALL_ZSH}" "${USERNAME}" "${USER_UID}" "${USER_GID}" "${UPGRADE_PACKAGES}" \
    && apt-get autoremove -y && apt-get clean -y && rm -rf /var/lib/apt/lists/* /tmp/library-scripts

devcontainer.json

{
    "name": "gpu-container",
    "build": {
        "dockerfile": "Dockerfile",
        "args": {
            "VARIANT": "2.4.0"
        }
    },
    "settings": {
        "terminal.integrated.shell.linux": "/bin/bash",
        "jupyter.alwaysTrustNotebooks": true,
        "http.proxySupport": "off"
    },
    "extensions": [
        "ms-python.vscode-pylance",
        "ms-python.python",
        "ms-toolsai.jupyter",
        "donjayamanne.python-extension-pack"
    ],
    "workspaceMount": "source=/home/remote/workspaces/gpu-container,target=/workspaces/gpu-container,type=bind,consistency=cached",
    "runArgs": ["--gpus", "all"],
    "remoteUser": "vscode"
}

devcontainer.jsonについて簡単に説明すると下記のような感じです。

• build.args.VARIANTでTensorFlowのバージョン指定が可能 (${VARIANT}-gpu-jupyterタグのイメージを取得)。
• extensionsにPython関連の拡張機能を設定。
• 設定jupyter.alwaysTrustNotebooksをtrueにし、notebookは全て信頼するように設定。
• workspaceMountで、WSL2上のフォルダをマウントし、データを永続化。
• runArgsで、GPUを使用するように設定。

上記で実行した結果は下記のスクリーンショットになります。notebookファイル(*.ipynb)を作成し、実行してみるとGPUを認識していることが分かります。

※VSCodeの拡張機能Python、Pylanceでnotebookを開くことが出来ます。jupyterではポート開放やセキュリティ(パスコード等)を考えないといけませんが、VSCodeだとローカルで気軽に開けるので個人的によく利用しています。またコンテナイメージ中に既にipykernelが入ってるので追加で操作しなくてもnotebookが利用できるようになっています。

各バージョンの動作確認とベンチマーク

TensorFlow benchmarksを利用してGPUの動作確認を行います。
下記のコマンドをDocker HubのTensorFlow中の該当するタグ名毎に実行し、また参考情報としてコンテナイメージ中のPython、及びCUDAのバージョン確認/整理した結果が下記の表になります。
※表中のベンチスコアは、実行後に表示されるtotal images/sec:の値を用いています。

console(GPU)

$ python tf_cnn_benchmarks.py --num_gpus=1 --batch_size=32 --model=resnet50 --variable_update=parameter_server

console(CPU)

$ python tf_cnn_benchmarks.py  --device=CPU  --data_format=NHWC --batch_size=32 --model=resnet50 --variable_update=parameter_server

No.	タグ名	python	TensorFlow	cuda	ベンチスコア	tf-bench ブランチ名
1	2.4.0-gpu-jupyter	3.6.9	2.4.0	11.0	105.37	master
2	2.3.1-gpu-jupyter	3.6.9	2.3.1	10.1	105.62	master
3	2.2.1-gpu-jupyter	3.6.9	2.2.1	10.1	104.53	master
4	2.1.2-gpu-jupyter	3.6.9	2.1.2	10.1	100.96	cnn_tf_v2.1_compatible
5	2.0.3-gpu-jupyter	3.6.9	2.0.3	10.0	100.15	cnn_tf_v2.0_compatible
6	1.15.4-gpu-jupyter	3.6.9	1.15.4	10.0	100.32	cnn_tf_v1.15_compatible
7	1.13.2-gpu-jupyter	2.7.15+	1.13.2	10.0	101.36	cnn_tf_v1.13_compatible
	…
8	1.5.0-gpu	2.7.12	1.5.0	9.0	93.62	cnn_tf_v1.5_compatible
9	(参考)ローカル	3.6.12	2.1.0	10.2	196.47	master
10	(参考)CPU(コンテナ)	3.6.9	2.4.0	11.0	8.14	master

結果は全てのTensorFlowバージョンで問題なく実行でき、またベンチスコアもそこまで変わらないという結果になりました。
コンテナをよく利用している人にとっては当たり前なのかもしれませんが、個人的にはGPUドライバ等のバージョン不一致でエラーが起きると思っていたので驚きました。
またこれまでローカルで環境構築する際は、TensorFlowのバージョンと、それが利用可能なドライバ、CUDAをいちいち調べてインストールしていたのと、利用するサンプルによってはTensorFlowのバージョンが異なっていて検証を断念することもあったので、気軽に変更できるのは非常に便利かと思います。

ただWindows上のローカル環境で実行した場合と比較して、スコアが約半分程度になってしまっています。
これについて、No.1のコンテナ環境と、No.9のローカル開発環境にて、バッチサイズ(--batch_size=)を変更した際、どの程度差があるのか比較した結果が下記の表になります。

バッチサイズ	スコア(コンテナ)	スコア(ローカル)	性能比
32	103.36	194.90	53.0%
48	138.54	201.36	68.8%
64	154.92	208.29	74.4%
80	169.99	208.12	81.7%
96	(エラー)	(エラー)	－

エラーになっている欄は、GPUメモリが確保できずに異常終了しているケースになります。
上記の通り、バッチサイズが大きくなればなるほどローカル環境との差が小さくなる結果となりました。
ハイパーバイザーの関係か不明ですが、メインメモリへのアクセスが多くなればなるほど遅くなるようなので、軽い処理ではCPUを利用した方が良さそうです。
またNVIDIAのページを見ると、パフォーマンスの改善に取り組んでいる最中で、APIの利用も現状では制限があるようなので、またアップデートされ改善が確認できたら記事を更新しようと思います。

終わりに

今回の検証結果で、開発環境もそうですが、運用環境でも異なるバージョンのTensorFlowやCUDAがお手軽に同居できるので、ますますコンテナが普及していきそうだと感じでいます。
その反面Anacondaが有償化したように、こうしたベースとなるコンテナイメージが有償化／ライセンス違反していかないか等注意深く見ていく必要がありそうです。
またリモート開発用のコンテナも同じ手順で手軽に準備できるので、環境構築作業等の時間短縮にもつながるかと思います。

関連記事

• WSL2 + 透過プロキシでプロキシ環境下のDockerを(少し)改善する - Qiita
• VSCode+WSL2+Dockerによるリモート開発環境の構築 (devcontainer.json概訳付き) - Qiita

参考記事

• tensorflow/tensorflow - Docker Hub
• tensorflow/benchmarks
• CUDA on WSL :: CUDA Toolkit Documentation
• microsoft/vscode-dev-containers
• 待ってました CUDA on WSL 2 - Qiita

商標

• NVIDIA、NVIDIAロゴ、GeForceは、米国およびその他の国におけるNVIDIA Corporationの商標または登録商標です。
• TensorFlowは，Google Inc.の米国およびその他の国における登録商標または商標です。
• Microsoft、Windows、Visual Studioは，米国Microsoft Corporationの米国およびその他の国における登録商標または商標です。
• Ubuntuは，Canonical Ltd.の商標または登録商標です。
• DockerおよびDockerのロゴはDocker,Inc.の米国およびその他の国における商標または登録商標です。
• その他記載の会社名、製品名、サービス名等はそれぞれの会社の商標または登録商標です。