PyTorchのDockerイメージを作成してみた


PyTorch v0.1.11のDockerイメージを作成して、examples/mnistを試してみました。

スポンサーリンク

PyTorch

バックエンドがTorchということだったので、KerasとTensorFlowの関係をイメージしていたのですが、dynamic neural networksという機能のためか、Chainerに似ている感じです。

Dockerイメージの作成

ここ(https://github.com/pytorch/pytorch#docker-image)を参照して、Dockerイメージを作成します。

$ cd ~/github
$ git clone https://github.com/pytorch/pytorch
$ cd pytorch
$ docker build -t pytorch-cudnnv6 .

出来上がったDockerイメージですが、サイズが結構大きく5.18GBとなっています。また、CUDA 8.0と、Dockerイメージ名の通りcuDNN 6.0に対応しています。

$ docker images pytorch-cudnnv6
REPOSITORY        TAG      IMAGE ID       CREATED      SIZE
pytorch-cudnnv6   latest   a113ed663fb9   2 days ago   5.18 GB

Dockerコンテナの起動

以下のコマンドで、一旦Dockerコンテナを起動してみましたが、

$ nvidia-docker run --rm -it --ipc=host \
    pytorch-cudnnv6

pytorch/examples(https://github.com/pytorch/examples)とpytorch/vision(https://github.com/pytorch/vision)を利用できるようにした方が良さそうです。

pytorch/examplesは、以下のコマンドでホスト側にcloneし、

$ cd ~/github
$ git clone https://github.com/pytorch/examples pytorch-examples

以下の-vオプションでDockerコンテナにマウントします。

$ nvidia-docker run -it --ipc=host \
    -v /home/sora/github/pytorch-examples:/workspace/examples \
    -v /home/sora/data/pytorch:/workspace/data \
    pytorch-cudnnv6

また、pytorch/visionは、pip install torchvisionでインストールする予定でしたが、pytorchが古いバージョンに書き戻されそうな感じだったので、以下のコマンドでインストールを行いました。

# cd /opt
# git clone https://github.com/pytorch/vision
# cd vision
# python setup.py install

今回は、Dockerコンテナにtorchvisionをインストールしましたが、Dockerイメージを再作成する際に、Dockerイメージそのものにインストールを行う予定です。

examples/mnistの実行

以下のコマンドで、mnistを実行します。なお、1回目の実行時に、MNISTデータセットが、/workspace/examples/data(ホスト側の~/github/pytorch-examples/data)にダウンロードされます。

# cd /workspace/examples/mnist
# python main.py
Train Epoch: 1 [0/60000 (0%)]	Loss: 2.390086
Train Epoch: 1 [640/60000 (1%)]	Loss: 2.350225
Train Epoch: 1 [1280/60000 (2%)]	Loss: 2.288934
...
Train Epoch: 10 [58240/60000 (97%)]	Loss: 0.104230
Train Epoch: 10 [58880/60000 (98%)]	Loss: 0.286794
Train Epoch: 10 [59520/60000 (99%)]	Loss: 0.076857

Test set: Average loss: 0.0546, Accuracy: 9813/10000 (98%)

まとめ

PyTorch v0.1.11のDockerイメージを作成して、examples/mnistを試してみました。