PyTorchの学習を可視化するには、
Facebookが公開しているvisdomを使うのが良さそうなのですが、
いろいろ試すのに、環境作ったり、立ち上げるのも手間なので、
jupyterとvisdomをまとめて動かせるdocker-imageを作りました。

本エントリは、その作成に関するメモです。

visdom
https://github.com/facebookresearch/visdom
※TensorFlowでいうTensorBoardみたいなものという理解で良さそうです。

Dockerfileと起動方法

最終的に作成したDockerfileと起動方法を示します。

Dockerfile

Dockerfileは以下の通りです。

FROM ubuntu:20.04

RUN apt-get update \
  && apt-get install -y python3 python3-pip supervisor
RUN pip3 install torch==1.7.1+cpu torchvision==0.8.2+cpu -f https://download.pytorch.org/whl/torch_stable.html \
  && pip3 install visdom jupyter
RUN echo "\
[supervisord]\n\
nodaemon=true\n\
[program:visdom]\n\
command=visdom\n\
[program:jupyter]\n\
command=jupyter notebook --ip=* --no-browser --allow-root\n\
stdout_logfile=/dev/stdout\n\
stdout_logfile_maxbytes=0\n\
stderr_logfile=/dev/stderr\n\
std_logfile_maxbytes=0\n\
" > /etc/supervisor/conf.d/supervisord.conf

EXPOSE 8888 8097
WORKDIR /work
CMD ["/usr/bin/supervisord"]

CUDA(nvidia-docker)を使いたい場合は、
ベースイメージ、インストールするPyTorchを変更します。
CUDA 11.0であれば、以下のような感じになると思います。

FROM nvidia/cuda:11.1-devel-ubuntu16.04

RUN pip3 install torch==1.7.1+cu110 torchvision==0.8.2+cu110 -f https://download.pytorch.org/whl/torch_stable.html \

ベースイメージ、PyTorchの指定は、
それぞれ、以下で調べます。

nvidia/cuda | dockerhub
https://hub.docker.com/r/nvidia/cuda/

GET STARTED | PyTorch
https://pytorch.org/get-started/locally/

イメージのビルド

以下のように、イメージをビルドします。

docker build . -t torchvisdom

起動方法

作成したイメージは、以下のポートでサービス起動するので、それぞれ指定して実行します。

jupyter: Port: 8888
visdom: Port: 8097

docker run -v `pwd`:/work -p 8888:8888 -p 8097:8097 torchvisdom

接続方法

ブラウザを開いて、以下のURLから各サービスを利用します。
jupyterはtokenの入力を求められるので、ターミナルに表示されているtokenを入力します。

jupyter: http://127.0.0.1:8888/
visdom: http://127.0.0.1:8097/

適当なサンプルを動かす

環境が出来たので、以下のページにあるサンプルに、
損失関数の可視化を追加して動かすことにします。

LEARNING PYTORCH WITH EXAMPLES | PyTorch
https://pytorch.org/tutorials/beginner/pytorch_with_examples.html

前段部分、学習部分のコードに分けていますが、
学習部分のコードを実行すると、
次の画像のように、visdomに損失関数の収束の様子が表示されます。

jupyter-visdom

前段部分のコードです。

import torch
import math
import visdom # visdomでの可視化のために追加
import numpy as np

dtype = torch.float
device = torch.device("cpu")
# device = torch.device("cuda:0") # Uncomment this to run on GPU

vis = visdom.Visdom(env='sample_1') # visdomでの可視化のために追加

# Create random input and output data
x = torch.linspace(-math.pi, math.pi, 2000, device=device, dtype=dtype)
y = torch.sin(x)

# Randomly initialize weights
a = torch.randn((), device=device, dtype=dtype)
b = torch.randn((), device=device, dtype=dtype)
c = torch.randn((), device=device, dtype=dtype)
d = torch.randn((), device=device, dtype=dtype)

学習部分のコードです。

learning_rate = 1e-6
for t in range(2000):
    # Forward pass: compute predicted y
    y_pred = a + b * x + c * x ** 2 + d * x ** 3

    # Compute and print loss
    loss = (y_pred - y).pow(2).sum().item()
    if t % 100 == 99:
        print(t, loss)
        vis.line(X=np.array([t]), Y=np.array([loss]), update='append', win='loss', name='loss') # visdomでの可視化のために追加

    # Backprop to compute gradients of a, b, c, d with respect to loss
    grad_y_pred = 2.0 * (y_pred - y)
    grad_a = grad_y_pred.sum()
    grad_b = (grad_y_pred * x).sum()
    grad_c = (grad_y_pred * x ** 2).sum()
    grad_d = (grad_y_pred * x ** 3).sum()

    # Update weights using gradient descent
    a -= learning_rate * grad_a
    b -= learning_rate * grad_b
    c -= learning_rate * grad_c
    d -= learning_rate * grad_d


print(f'Result: y = {a.item()} + {b.item()} x + {c.item()} x^2 + {d.item()} x^3')

参考にした情報

Dockerで、複数プロセス動かす方法として、
Supervisorを使う方法が紹介されています。

Supervisor を Docker で使う | docker-docs-ja https://docs.docker.jp/engine/admin/using_supervisord.html

Supervisorを使うと、
jupyterのtokenが標準エラー出力に表示されなくなってしまうので、
次の方法で変更しています。

[メモ] Dockerでsupervisordを使うときにログを標準出力する | sugiki / Qiita https://qiita.com/sugiki/items/86515393664cb5ce16ff

以上。