PyTorch の Define by Run を可視化

Deep Learningフレームワークにおける計算グラフ

TensorFlow, PyTorchなどのDeep Learningフレームワークにおける計算グラフの構築方法が，大きく分けて2パターンに分類される。1つ目は Define by Run 型で，2つ目は Define and Run 型である。

それぞれの特徴

Define by Run型

• Chainer, PyTorchがこちらの方式を採用している
• データを流すときに，計算グラフの構造を決定
• 動的な計算グラフを定義することが可能

Define and Run型

• TensorFlow(v1)がこの方式を採用している
• 計算グラフの構造を予め定義しておき，そこにデータを流す

Define by Run が動的な計算グラフを構築できることを確認

概要

• 入力データによって計算グラフが変化することを確認
• 本稿では，Recurrent Neural Net (RNN) を可視化
• torch_vizのmake_dot()を用いて，系列長が2と3の場合の計算グラフを出力

結果

• 系列長2の場合 (左)
  • ループに対する重みとバイアス(fc_h.weight と hc_h.bias)が1度しか使用されない
• 系列長3の場合 (右)
  • ループする回数が増えたため，2度使用されていることがわかる

ソースコード

可視化

visualize_RNN.py

import torch

from torchviz import make_dot
from RNN import RNN

INPUT_SIZE = 1
TIME_LENGTH = 3
# TIME_LENGTH = 2

model = RNN()
data = torch.randn(1, TIME_LENGTH, INPUT_SIZE)

y = model(data)

image = make_dot(y, params=dict(model.named_parameters()))
image.format = "png"
image.render(f"RNN_{TIME_LENGTH}s")

ネットワークの定義

RNN.py

import torch
import torch.nn as nn

class RNN(nn.Module):
    def __init__(self):
        super(RNN, self).__init__()
        # 入力に対する全結合層
        self.fc_in = nn.Linear(1, 64)
        # ループに対する全結合層
        self.fc_h = nn.Linear(64,64)
        # 出力用の全結合層
        self.fc = nn.Linear(64, 10)

    def forward(self, x):
        _, length, _ = x.size()
        h = self.fc_in(x)
        h = torch.tanh(h)

        # 入力データの系列長に合わせてループする
        for i in range(1, length):
            h = self.fc_in(x[:, i, :]) + self.fc_h(h)
            h = torch.tanh(h)

        h = self.fc(h)

        return h