Windows11上でCUDA + Pytorch + Jupyter環境をdocker-compose.yamlで構築する

はじめに

Windows上で手軽に構築できて環境を汚染しないJupyterの環境がほしいと思ったことことがきっかけで調べました。
本稿で使っているPytorchのバージョンは少し前のものだと思いますので、有識者は以下のサイトでほしいバージョンを探してください。
https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch

参考記事

以下の記事で使われているコマンドをもとにdocker-compose.yamlを作成しました。
https://qiita.com/taro_kawasaki/items/6738ada8e1f9ec0ed24a
https://qiita.com/ao_log/items/5438f2aaf5c2b46d2ccb

手順

準備

Docker DesktopとWSL2のインストール

WSL2とDocker Desktopをインストールし設定してください。
好みのやり方にもよりますが、私は以下の「Docker Desktopのインストールとその WSL 2 バックエンドを使う方法」を取りました。設定が楽でした。
https://qiita.com/ksasaki/items/ee864abd74f95fea1efa

作業フォルダの作成

次に、任意の場所に「workspace」フォルダを作成します。

docker-compose.yamlの記述

以下をもとにdocker-compose.yamlファイルを作り、どこか分かりやすいところに保存しておいてください。
ただし1箇所、sercives->jupyter->volumesの左側の「E:/workspace」を、先ほど作ったworkspaceフォルダのパスに変えてください。「：」から右側はそのままで大丈夫です。

docker-compose.yaml

version: '3.8'

services:
  jupyter:
    image: nvcr.io/nvidia/pytorch:21.07-py3
    container_name: jupyter
    ports:
      - "9000:8888"
    volumes:
      - "E:/workspace:/workspace" # ← 環境に応じて変更してほしい行はここです
    working_dir: /workspace
    command: jupyter-lab --no-browser --port=8888 --ip=0.0.0.0 --allow-root --NotebookApp.token=''
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

動作確認

ターミナルまたはPowerShellを開き、上記のdocker-compose.yamlファイルを保存した階層まで移動します。
そこで、下記コマンドを実行してください。

$ docker-compose up

初回起動時はイメージの取得で時間を要しますが、それが終わればエラーでも表示されていない限り起動していると思いますので、 http://localhost:9000/lab でJupyterLabを開けます。

新しいノートブックを開き、以下のコマンドで確認できます。

import torch

torch.__version__

torch.cuda.is_available()

こんな感じでバージョンとCUDAが利用可能であることが示されれば完了です。

一度構築すれば、次回以降はDocker Desktopからも立ち上げられます。