🐳

WSL2とDocker Desktopでtensorflow-gpu 1.X系の環境構築をする

2022/08/24に公開

Docker

普段はtensorflow 2.X系を使っているのですが、どうしてもtensorflow 1.X系を使わないといけない状況になり、Docker Desktop(WSL2バックエンド)で環境構築したのでその手順を記載します。

①WindowsにGPUドライバをインストール

まず、Windowsにドライバ(CUDA)のインストールを行います。
既にCUDAをWindowsにインストール済みの方はこの手順を飛ばして頂いて構いません。

ちなみに、NVIDIAの公式のドキュメントに記載がありますが、最新のドライバーはWSL2を完全サポートしており、Windows NVIDIA GPUドライバーがインストールされると、WSL2内でも自動的にCUDAが使用可能になります。
したがって、WSL2上で改めてドライバのダウンロードをする必要はありません。

Dockerイメージの取得

Docker Desktopのインストールがまだな方は下記からインストールして下さい。
公式の手順通りにインストールを行えば特に問題ありません。

TensorFlowのDockerイメージはDocker Hubで取得可能です。

今回はtensorflow-gpu 1.15.5のイメージを取得します。

Windows側のターミナルで下記を実行します。

docker pull tensorflow/tensorflow:1.15.5-gpu

コンテナの起動

Dockerイメージの取得が完了したら、コンテナの起動を行います。
オプション指定はお好みでOKですが、GPUを認識させるために「--gpus all」の指定が必要になります。

docker run --name tensorflow-gpu-1.15.5 -it --gpus all tensorflow/tensorflow:1.15.5-gpu

GPUを認識できているかチェック

下記を実行し、GPUを認識できているか確認します。

gpu-test.py

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

次のような表示が出ていればGPUの認識ができています。

incarnation: 10845244744991911999
physical_device_desc: "device: XLA_GPU device"
, name: "/device:GPU:0"
device_type: "GPU"
memory_limit: 7128429364
locality {
  bus_id: 1
  links {
  }
}

一応、CUDAのバージョンも確認しておきます。

$ nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2018 NVIDIA Corporation
Built on Sat_Aug_25_21:08:01_CDT_2018
Cuda compilation tools, release 10.0, V10.0.130

機械学習を実行してみる

scikit-learnやkerasをインストールして、irisデータセットを分類する簡単なNNモデルの作成・学習・テストをしてみます。

nn-test.py

from sklearn import datasets
from sklearn import preprocessing
from keras.utils import np_utils
from sklearn.model_selection import train_test_split
from keras.models import Sequential
from keras.layers.core import Dense, Activation

iris = datasets.load_iris()
X = iris.data
y = iris.target

X = preprocessing.scale(X)
y = np_utils.to_categorical(y)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0)

model = Sequential()
model.add(Dense(100, input_shape=(4,)))
model.add(Activation('relu'))
model.add(Dense(100))
model.add(Activation('relu'))
model.add(Dense(3))
model.add(Activation('softmax'))
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

model.fit(X_train, y_train, epochs=50, batch_size=1, verbose=1)

test_loss, test_acc = model.evaluate(X_test, y_test, batch_size=1, verbose=0)

print('accuracy:', test_acc)

実行結果

（省略）
Epoch 50/50
105/105 [==============================] - 1s 5ms/step - loss: 0.0123 - accuracy: 1.0000
accuracy: 0.9777777791023254

CPUと比較し、爆速で学習が進むことが確認できました。
お疲れ様でした。

終わりに

Windows上でtensorflow-gpu 1.X環境を構築しようとすると、tensorflow-gpuのバージョンに応じたCUDA ツールキットやcuDNNをインストールする必要があって面倒ですが、Dockerを使うと簡単にセットアップすることができます。
是非試してみて下さい。