M1 macな人たちに贈るpytorch+jupyterlab docker image

はじめに

(2023/02/28 追記あり)

PCを変えるたびに開発環境やらtex環境やらを構築するのが面倒なので、当方はなるべくdockerイメージを使って環境構築するようにしています。

pythonを使うときは多くの場合にデータ分析やら機械学習を行うので、docker hubにあるpytorchイメージにjupyterlabやらscikit-learnやらを導入したイメージを使っていました。しかし、大変残念なことに、docker hubにあるpytorchイメージはlinux/amd64版のみが提供されており、linux/arm64（aarch64）なM1 mac民は使うことができません。

M1 macではまだnvidia GPUが使えないことから、心機一転してCPU実行のみをサポートする小さめのイメージを作成しました。

検証環境

1. Docker version 20.10.2, build 2291f61 running on Ubuntu 18.04.05 LTS x86_64
2. Docker version 20.10.7, build f0df350 running on macOS 10.15.7 x86_64
3. Docker version 20.10.7, build f0df350 running on macOS 11.4 aarch64

使い方

拙作のpytorch+jupyterlabをpullして実行してください。
latestタグのイメージを作っていないため、利用可能なタグを確認して実行してください。

(2023/02/28追記) いつからか記憶してませんが、latestタグを付けました。

CPUしか使わない人

本記事執筆時点ではpytorchはversion 1.9.0が最新です。ですので、以下のように実行します。
初回実行時に自動的にpullされると思います。

% docker run --rm -d -v $PWD:/app -p 8888:8888 \
  pman0214/pytorch_jupyterlab:1.9.0

(2023/02/28追記) バージョン1.13.1を追加してあります。latestタグも作成してあるので、以下でpullすれば最新版になります。

% docker run --rm -d -v $PWD:/app -p 8888:8888 \
  pman0214/pytorch_jupyterlab

あとはブラウザでhttp://localhost:8888にアクセスすればjupyterlabが使えます。
jupyterlabのパスワードはjupyterとしてあります。

GPUも使う人

linux/amd64な環境でnvidia GPUを使う場合は、GPUサポートが付いているタグを指定して実行してください。
なお、GPUを使うためには--gpus=allなどのオプションを付ける必要があります。dockerもバージョン19.03以降である必要があると思います（未検証）。

% docker run --rm -d -v $PWD:/app -p 8888:8888 --gpus=all \
  pman0214/pytorch_jupyterlab:1.9.0-cuda11.1-cudnn8-devel

(2023/02/28追記) GPU版もバージョン1.13.1があります。latest-gpuタグでGPU対応の最新版をpullできます。arm64版もありますが、arm64+nvidia GPUという環境では動かないニセモノです（ミスで作成されたものですので無視してください）。

% docker run --rm -d -v $PWD:/app -p 8888:8888 --gpus=all \
  pman0214/pytorch_jupyterlab:latest-gpu

あとはブラウザでhttp://localhost:8888にアクセスすればjupyterlabが使えます。
jupyterlabのパスワードはjupyterです。

GPUにアクセスできているかどうかは、jupyterlab上で以下のようにpythonコードを実行すれば確認できます。

import torch
print(torch.cuda.is_available())

TrueならGPUが使用可能で、GPU情報などを取得できます。

print(torch.cuda.get_device_name())

EPS画像出力への対応

論文などではいまだにEPS画像を提出せよなどという場合もあるため、EPS出力をサポートしたイメージも作りました。
使い方は上記のものと同じで、イメージ名だけが異なります。

ビルドしたい人向け

ビルドに用いたファイル群はGitHubで公開しています。どうせならと思い、linux/amd64もlinux/arm64もサポートするmulti-architecture buildをしています。
なお、GPU対応版は別のベースイメージを使ってlinux/amd64用のみをビルドしました。

docker contextで接続先マシンを切り替えてbuildする手法もありますが（こっちのが速いらしい、当たり前か）、ここではより一般に適用可能なQEMUを使う場合で説明を行います。
ビルド用にQEMUを使ったコンテナが作成され、その中でビルドが実行されます。ビルドにはBuildkitが使われるので、ビルド用のコンテナなどを自分で準備する必要はありません。

ビルド方法の詳細はdockerの公式ドキュメントを参照するとよいです。概ね同じ手順になっていると思います。

準備（linuxでビルドする場合のみ）

macやwindowsでビルドする場合は、最新版のDocker desktopが入っていれば準備は不要（になったはず）です。

experimental featureの有効化

まず、~/.docker/config.jsonに以下のような記述を追加するなどして、experimental featureを有効化します。

{
   ...
   "experimental": "true"
   ...
}

うまく有効化されれば、buildxコマンドが使えるようになります。

% docker buildx ls
NAME/NODE DRIVER/ENDPOINT STATUS  PLATFORMS
default * docker
  default default         running linux/amd64, linux/386

有効化できていない場合には、このようになります。確認していませんが、configを書いた後にdockerデーモンを再起動するなどが必要かもしれません。

% docker buildx ls
docker: 'buildx' is not a docker command.
See 'docker --help'

QEMUの導入

ホストOSの機能で導入するとロクなことにならないので（経験者談）、dockerで配布されている導入キットを利用します。なお、このキットはBuildkitにも導入されています。

念のためインストールされている（かもしれない）QEMUを削除してから導入します。

% docker run --privileged --rm tonistiigi/binfmt --uninstall "qemu-*"
% docker run --privileged --rm tonistiigi/binfmt --install all

ビルダの準備

まず、buildxコマンドでビルダを準備します。名前は任意です。ここではmultiarchとしています。
--driverオプションでdocker-containerを指定することで、QEMUを実行するdocker-container上でビルドを行う指示をしています。

% docker buildx create --name multiarch --driver docker-container

これで、buildx lsの出力はこのようになるはずです。

% docker buildx ls
NAME/NODE    DRIVER/ENDPOINT             STATUS   PLATFORMS
multiarch    docker-container
  multiarch0 unix:///var/run/docker.sock inactive
default *    docker
  default    default                     running  linux/amd64, linux/386, linux/arm64, linux/riscv64, linux/ppc64le, linux/s390x, linux/arm/v7, linux/arm/v6

*が付いているものが、使用されるビルダ（現在選択されているビルダ）です。
作成したビルダmultiarchを使用するように変更します。

% docker buildx use multiarch

ビルダが切り替わっていることが確認できます。multi-architecture buildしないときはdocker buildx use defaultとしてdefaultに戻してください。

% docker buildx ls
NAME/NODE    DRIVER/ENDPOINT             STATUS   PLATFORMS
multiarch *  docker-container
  multiarch0 unix:///var/run/docker.sock inactive
default      docker
  default    default                     running  linux/amd64, linux/386, linux/arm64, linux/riscv64, linux/ppc64le, linux/s390x, linux/arm/v7, linux/arm/v6

ビルダのコンテナを作成します（実はこれをやらなくても、自動で作成されるという噂も・・・）。

% docker buildx inspect --bootstrap

ビルド

アーキテクチャを指定して、buildxコマンドを使ってビルドします。
注意するのは、--loadを付ける点です。付けないと、ビルドした結果は破棄されます（という警告も表示されます）。

--loadの代わりに--pushとすればdocker hubにpushされます。この場合にはあらかじめdocker loginしておき、-tオプションで適切なタグを指定しておいてください。

% git clone https://github.com/pman0214/docker_pytorch-jupyterlab.git
% cd docker_pytorch-jupyterlab
% docker buildx build \
  --platform linux/amd64,linux/arm64 \
  -t "pytorch_jupyterlab_latex:1.9.0" \
  --build-arg VER=1.9.0 . \
  --load

複数のアーキテクチャのビルドは並行して実行されます。
なお、QEMUを用いてビルドするのでホストと異なるアーキテクチャ上でのビルドは特に遅いです。例えば、Core i9 3.6GHz 16core with 64GB memoryなUbuntu 18.04マシンで、上記ビルドには全部で2時間ほど要しました。

GPU版のビルド

GPU版はamd64なマシンがあれば、multi-architecture buildをする必要はありません。
Dockerfileが異なるので、その指定だけ忘れずにビルドしてください。

% docker build \
  -f Dockerfile_cuda \
  -t "pytorch_jupyterlab_latex:1.9.0-cuda11.1-cudnn8-devel" \
  .

おわりに

本記事では、pytorch + jupyterlabなdockerイメージの使い方及び作り方を説明しました。これでM1 macな人たちもdockerさえ入っていればaarch64ネイティブで実行されるpytorchライフを送れますね。

ビルドをgithub action化するとさらに幸せになれそうですが、未着手です。