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

# はじめに
(2023/02/28 追記あり)

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

pythonを使うときは多くの場合にデータ分析やら機械学習を行うので、[docker hubにあるpytorchイメージ](https://hub.docker.com/r/pytorch/pytorch)にjupyterlabやらscikit-learnやらを導入したイメージを使っていました。しかし、大変残念なことに、[docker hubにあるpytorchイメージ](https://hub.docker.com/r/pytorch/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
1. Docker version 20.10.7, build f0df350 running on macOS 10.15.7 x86_64
1. Docker version 20.10.7, build f0df350 running on macOS 11.4 aarch64

# 使い方

拙作の[pytorch+jupyterlab](https://hub.docker.com/r/pman0214/pytorch_jupyterlab/)をpullして実行してください。
`latest`タグのイメージを作っていないため、[利用可能なタグ](https://hub.docker.com/r/pman0214/pytorch_jupyterlab/tags)を確認して実行してください。

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

## CPUしか使わない人

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

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

**(2023/02/28追記)** バージョン1.13.1を追加してあります。`latest`タグも作成してあるので、以下でpullすれば最新版になります。
```shell
% docker run --rm -d -v $PWD:/app -p 8888:8888 \
  pman0214/pytorch_jupyterlab
```

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

## GPUも使う人

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

```shell
% 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という環境では動かないニセモノ**です（ミスで作成されたものですので無視してください）。
```shell
% docker run --rm -d -v $PWD:/app -p 8888:8888 --gpus=all \
  pman0214/pytorch_jupyterlab:latest-gpu
```

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

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

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

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

```python
print(torch.cuda.get_device_name())
```

# EPS画像出力への対応

論文などではいまだにEPS画像を提出せよなどという場合もあるため、[EPS出力をサポートしたイメージ](https://hub.docker.com/r/pman0214/pytorch_jupyterlab_latex/)も作りました。
使い方は上記のものと同じで、イメージ名だけが異なります。

# ビルドしたい人向け

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

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

ビルド方法の詳細は[dockerの公式ドキュメント](https://docs.docker.com/desktop/multi-arch/)を参照するとよいです。概ね同じ手順になっていると思います。

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

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

### experimental featureの有効化

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

```json
{
   ...
   "experimental": "true"
   ...
}
```

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

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

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

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

### QEMUの導入

ホストOSの機能で導入するとロクなことにならないので（経験者談）、[dockerで配布されている導入キット](https://github.com/tonistiigi/binfmt/)を利用します。なお、このキットは[Buildkit](https://github.com/moby/buildkit)にも導入されています。

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

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

### ビルダの準備

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

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

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

```shell
% 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を使用するように変更します。

```shell
% docker buildx use multiarch
```

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

```shell
% 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
```

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

```shell
% docker buildx inspect --bootstrap
```

## ビルド

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

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

```shell
% 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`が異なるので、その指定だけ忘れずにビルドしてください。

```shell
% 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化するとさらに幸せになれそうですが、未着手です。