dockerでとにかくGPGPU環境が欲しいときのワンライナー（Tensorflow, Pytorch, Keras諸々対応）

はじめに

この記事はdocker、NVIDIAの環境がそろった（nvidia-docker nvidia-smi などのコマンドが通る）環境を想定して、はじめます。

ワンライナー

nvidia-docker run -it  --name your_env_name -v /home/your_home_dir:/home/your_home_dir --env CUDA_​VISIBLE_​DEVICES=0 nvcr.io/nvidia/tensorflow:19.05-py3 bash

説明

• your_env_name: つくるdockerの仮想環境名
• your_home_dir: ホストOS上でのあなたのホームディレクトリ。-vオプションで仮想環境と共有することができる。ホームを共有すると便利
• CUDA_VISIBLE_DEVICES: 仮想環境で使うGPUのIDを指定する。上では0のみを使用する。通常、0~(GPU個数-1)までのGPUが指定可能。複数個指定する場合は,で区切る。
• nvcr.io/nvidia/tensorflow:19.05-py3: tensorflowを使用する場合の2019/6/14現在の最新版参照環境。以下のURL先を参照し適宜変更
  • 参照環境用URL
    • https://ngc.nvidia.com/catalog/containers
    • https://hub.docker.com/

これでとりあえず上のワンライナーでtensorflow環境のBASHに入ることができる。