PyTorchで書いたモデルをTensorBoardでPNGに吐き出す

なにをしたいか

• PyTorchで書いたモデルを、簡単に画像に出力したい。
• graphvizを直接使うのは面倒なので避ける
• Example Codeを実行してみる

環境

• torch==1.1.0.post2
• tensorboard==1.14.0
• future==0.17.1

Example Code

resnet50をMNISTデータでトレーニングする場合(example codeをコピペ)

モデルをTensorBoard用に出力

import torch
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms

# Writer will output to ./runs/ directory by default
writer = SummaryWriter()

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
trainset = datasets.MNIST('mnist_train', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
model = torchvision.models.resnet50(False)
# Have ResNet model take in grayscale rather than RGB
model.conv1 = torch.nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
images, labels = next(iter(trainloader))

grid = torchvision.utils.make_grid(images)
writer.add_image('images', grid, 0)
writer.add_graph(model, images)
writer.close()

tensorboardを起動

tensorboard --logdir=runs

自分でもやってみる

TorchVisionのAlexNetをPNGに吐き出してみる
(モデルを書き出すだけなのでデータセットは用意しない)

TorchvisionのAlexNetをTensorBoardに渡す

import torch
from torchvision import models
from torch.utils.tensorboard import SummaryWriter

# Writer will output to ./runs/ directory by default
writer = SummaryWriter()

image = torch.randn(1,3,224,224)
model = models.AlexNet()
writer.add_graph(model, image)
writer.close()

TensorBoardを起動

tensorboard --logdir=runs

PNGに保存

TensorBoardをブラウザから表示した上で、「Download PNG」ボタンをクリック

ちょっと気になる点

• インタラクティブな表示とはいえ、表示の仕方があまりかっこよくない気がします（小並感）