DockerからtensorboardXをつかう

tensorboardXは,Tensorflowの可視化ライブラリであるtensorboardをPyTorchから使用できるようにしたパッケージです。こちらの環境ではdocker環境からjupyter notebookを使っていましたが、tensorboardXも使えるようにした際のメモです。

行ったこと

1. tensorboardx も使えるようなdocker環境を修正した。
2. tensorboardXの入力となるグラフファイルを作成し、tensorboardXの基本動作を確認した。

動作環境

OS

$ cat /etc/lsb-release 
DISTRIB_ID=Ubuntu
DISTRIB_RELEASE=16.04
DISTRIB_CODENAME=xenial
DISTRIB_DESCRIPTION="Ubuntu 16.04.6 LTS"

dockerの環境

こちらで用いている環境をベースにしています。

1. tensorboardX も使えるようにdocker環境を修正

tensorboardXでは、ニューラルネットワークの共通フォーマットであるONNX(Open Neural Network Exchange)形式を可視化するので、まずはPytorchで作成されたモデルをONNX形式に変換する必要があります。
tensorboardXを使えるようにするためには、tensorflowとtensorboardxをインストールする必要があるため、Dockerfileに追加しました。

Dockerfile

FROM nvcr.io/nvidia/pytorch:20.03-py3

ARG DEBIAN_FRONTEND=noninteractive

RUN apt-get update && \
    apt-get install -y python3-tk && \                      ## GUIに関するライブラリ
    pip install --upgrade pip && \
    apt-get install -y libsm6 libxext6 libxrender-dev && \　## OpenCVに必要なライブラリ
    pip install opencv-python && \
    pip install pytest && \　　　　　　　　　　　　　　　　　　　 ## テスト環境に必要
    pip install nose && \　　　　　　　　　　　　　　　　　　　　 ## テスト環境に必要
    pip install ipywidgets && \　　　　　　　　　　　　　　　　　## UI
    pip install tensorflow --ignore-installed --user && \　
    pip install tensorboardx

RUN echo 'export QT_X11_NO_MITSHM=1' >> ~/.bashrc && \
    source ~/.bashrc

上記Dockerfileを用いて、docker imageを作成します。

$cd **/work/  [Dockerfileのあるフォルダ]
$docker build -t [repository名]:[tag名] . 
例）$docker build -t test:1 . 

作成したimageを用いてdockerのcontainerを立ち上げます。
(web cameraの設定はこちらを参照ください。)

$xhost +
$docker run --gpus all -it --rm \
           --device /dev/video0:/dev/video0:mwr \   #web camera
           -p 10000:8888 \                          #jupyter用のport forwarding
           -p 6006:6006 \                           #tensorboardx用のport forwarding
           -v /[working directory path]/:/work \
           -v /tmp/.X11-unix:/tmp/.X11-unix \
           -e DISPLAY=$DISPLAY test:1

$cd /work

コンテナが起動したら、tensorboardXを立ち上げる準備ができました。

2. グラフファイルの作成とtensorboardXの基本動作

可視化したいネットワークモデルとして、pytorchに入っているvgg19のモデルファイルを作成します。

graph

import torch
import torchvision

net = torchvision.models.vgg19(pretrained=True)
net.train()

vgg19

VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (17): ReLU(inplace=True)
    (18): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (19): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (24): ReLU(inplace=True)
    (25): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (26): ReLU(inplace=True)
    (27): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (31): ReLU(inplace=True)
    (32): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (33): ReLU(inplace=True)
    (34): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (35): ReLU(inplace=True)
    (36): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

次に、上記モデルをグラフファイルに出力する。グラフ以外にも様々なものが出力できる(参考資料)。

from tensorboardX import SummaryWriter

batch_size = 5
dummy_input = torch.rand(batch_size, 3, 368, 368)

with SummaryWriter("./tbx/", comment='test example vgg19') as w:
    w.add_graph(net, (dummy_input, ), True)

こうしてカレントディレクトリの下にtbxというフォルダが作成され、その中にグラフファイルevents.out.tfevents.*** というファイルが生成されます。1章で起動したdocker内で、tensorboardを立ち上げます。ここで、グラフファイルがあるフォルダまでのパスを以下のように指定します。

#tensorboard --logdir="./tbx/" --bind_all

ホスト側でブラウザを立ち上げ, URLとして http://localhost:6006 を入力すると、ブラウザ上にtensorboardxの画面が立ち上がります。

グラフ内のブロックをダブルクリックすると、より詳細な情報が提示されます。

最後に、Tensorbordxを終了するには、ターミナルでCtrl+Cを入力します。

参考にさせていただいた記事

1.[tensorboardx] (https://tensorboardx.readthedocs.io/en/latest/index.html)
2.graph_addのサンプルコード
3.PyTorchによる発展ディープラーニング
4.Dockerでポートフォワーディング解説