DockerでPyTorchの実行環境を整備する 2019年11月 #Python

概要

PyTorchによるDeep Learningの実行環境をDockerコンテナで整備する
Docker version19.03以降，ホストOSのGPUドライバとnvidia-container-runtimeさえあれば他はコンテナに閉じ込められる
これで大体の場合コマンド一発で環境を再現できる

準備

ざっくりと以下の用意が必要．

dockerのversion19.03以降
- 公式ドキュメントを参照(このリンクはUbuntuの場合のドキュメント)
nvidia-container-runtime
- GitHubのREADME.mdを参考にインストール
NVIDIA Driver
- 公式サイトを参考に自分のGPUに合わせたドライバをインストール

Dockerfileを作成する

コンテナにはCUDAとCuDNNが入り，ホスト側のGPU Driverを利用する形になる．

このDockerfileではUbuntu:18.04にCUDA:10.1とCuDNN:7が入ったイメージをベースにPython 3.7.1とPythonのパッケージをインストールしている．

FROM nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04

WORKDIR /code
ENV PYTHON_VERSION 3.7.1
ENV HOME /root
ENV PYTHON_ROOT $HOME/local/python-$PYTHON_VERSION
ENV PATH $PYTHON_ROOT/bin:$PATH
ENV PYENV_ROOT $HOME/.pyenv
ENV TZ=Asia/Tokyo
ADD requirements.txt /code
RUN export DEBIAN_FRONTEND=noninteractive && \
    apt-get -y update && \
    apt-get -y upgrade && \
    apt-get -y install tzdata
RUN apt-get -y install git make build-essential libssl-dev zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev && \
     apt-get -y install wget curl llvm libncurses5-dev libncursesw5-dev xz-utils tk-dev libffi-dev liblzma-dev && \
    git clone https://github.com/pyenv/pyenv.git $PYENV_ROOT && \
    $PYENV_ROOT/plugins/python-build/install.sh && \
    /usr/local/bin/python-build -v $PYTHON_VERSION $PYTHON_ROOT && \
    rm -rf $PYENV_ROOT && \
    pip install -r requirements.txt

requirements.txtの中身の例

requirements.txt

torch
torchvision

コンテナ環境を利用する

まずはコンテナをビルドする．
Dockerfileを作成したディレクトリで以下のコマンドを実行する．

docker build . -t torch3.7

コンテナに入りGPUを認識できているかを確認する．
nvidia-smiコマンドで以下の表示が出ればok．

$ docker run -it --gpus all torch3.7 bash
$ nvidia-smi
 Tue Nov 19 15:01:12 2019       
 +-----------------------------------------------------------------------------+
 | NVIDIA-SMI 430.50       Driver Version: 430.50       CUDA Version: 10.1     |
 |-------------------------------+----------------------+----------------------+
 | GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
 | Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
 |===============================+======================+======================|
 |   0  GeForce GTX 166...  Off  | 00000000:1F:00.0 Off |                  N/A |
 | 45%   36C    P0     1W / 120W |      0MiB /  5944MiB |      0%      Default |
 +-------------------------------+----------------------+----------------------+

 +-----------------------------------------------------------------------------+
 | Processes:                                                       GPU Memory |
 |  GPU       PID   Type   Process name                             Usage      |
 |=============================================================================|
 |  No running processes found                                                 |
 +-----------------------------------------------------------------------------+

最後にPyTorchのtorch.TensorをGPU上のメモリに載せられることを確認する．
torch.TensorをGPU上で利用できた．

$ python
>>> import torch
>>> torch.cuda.is_available()
True
>>> x = torch.Tensor([0, 1])
>>> x
tensor([0., 1.])
>>> x.cuda()
tensor([0., 1.], device='cuda:0')

ちなみに...

ホストとコンテナでvolumeを同期して使うことが多くなると思われるのでその場合は以下のコマンドを適宜変更して利用すればok．

docker run -it -v /path/to/your_project_directory:/code --gpus all torch3.7 bash