Dockerでtensorflow-gpuとかを使えるようにする

既に走らせてるコンテナであ〜やっぱりgpu使いたいな〜ってとき

  • 必要なもの:予めPCにはnvidia driverをインストールしてnvidia-smiができるようにしておく

コンテナをイメージに保存

コンテナからイメージを作成して一旦保存する

まずコンテナが起動していたら以下のコマンドで停止する

$ docker stop [対象のコンテナ名]

次に以下のコマンドでイメージ化する

$ docker commit [対象のコンテナ名] [保存するイメージ名]

保存したイメージはdocker image listで確認できる

また、gpu対応させるときのコンテナ名を同じ名前にしたいときはコンテナ名を変えておく

$ docker rename [対象のコンテナ名] [保存するイメージ名]

GPU対応のコンテナを新たに作成

今のDockerではdocker runするときに--gpus allオプションを加えるとgpuが使えるようになる

なので先ほど作ったイメージから--gpus allオプションをつけてコンテナを作る

$ docker run -it --gpus all --name [新規作成するコンテナ名] [さっき作ったイメージ名] bash

起動したコンテナ上でnvidia-smiコマンドが成功したらOK

依存関係をまとめてtensorflow-gpuインストール

condaを使うとcudaとかcudnnとかの依存関係もまとめてインストールしてくれるので楽

$ conda install -y tensorflow-gpu

以下のpythonコードが動いて学習中にGPU使用率が上がってれば成功

import tensorflow as tf
mnist = tf.keras.datasets.mnist

(x_train, y_train),(x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(512, activation=tf.nn.relu),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation=tf.nn.softmax)
])
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(x_train, y_train, epochs=10)
model.evaluate(x_test, y_test)