More than 3 years have passed since last update.

Dockerを使ってpytorchでcuda11.0を使用する

Last updated at 2020-10-14Posted at 2020-10-14

2020年10月14日現在では、pytorch公式にはcuda11.0向けのインストール方法が掲載されていません。
そこで、Dockerを使用したcuda11.0向けのpytorchの環境構築の手順を説明します。

前提条件

前提としてDockerとnvidiaドライバーがインストールされているものとします。
これらのインストールは他の記事などを参照してください。

$ nvidia-smi

上記のコマンドで、cudaのバージョンが11.0になっていることも合わせて確認してください。

nvidiaのサイトにアクセスしてユーザー情報を登録後、Pytorchの画面を開いてください。

$ sudo docker pull nvcr.io/nvidia/pytorch:20.09-py3

pull commandに上記のコマンドがあるのでこれをコピーしてターミナルで実行します。

$ sudo docker images

このコマンドでインストールできたかを確認してください。
nvcr.io/nvidia/pytorch 20.09-py3
というimageがあれば成功です。

先ほどのimageを一覧表示するコマンドを実行して、pytorchのIMAGE IDを確認してコピーして下記コマンドを実行します。

$ sudo docker run -it --gpus all [IMAGE ID] bash

その後、bashで

$ python
>> import torch
>> print(torch.cuda.is_available())

これでtrueが返ってきたら成功です。
bashはcontrol + dで抜けれます。(Linux)

一度、bashから抜けてください。

$ sudo docker ps -a

上記のコマンドで先ほどbashを実行した際に作成されたcontainerが確認できます。これのCONTAINER IDを使用します。

$ sudo docker inspect [CONTAINER ID]

control + fでWorkingDirを探します。
さらに、自分が作業を行いたいPC上のフォルダのパスをコマンドで確認します。

$ pwd

その後、先ほどのnvidiaのサイトにoverviewで説明されているようにコンテナを作成します。

$ sudo docker run --gpus all -it --rm -v [作業を行いたいフォルダのパス]:[workingDir] nvcr.io/nvidia/pytorch:20.09-py3

私のケースだと

$ sudo docker run --gpus all -it --rm -v /home/myname/Desktop/hoge:/workspace　nvcr.io/nvidia/pytorch:20.09-py3

その後、containerが立ち上がることが確認できるはずです。

タイポなどがある可能性があるので気をつけてください。
公式がインストール方法を掲載してくれるまでの繋ぎになれば幸いです。