如何在 Pytorch 中可视化网络？

以下是使用不同工具的三种不同图形可视化。

为了生成示例可视化，我将使用一个简单的 RNN 来执行来自情绪分析：

class RNN(nn.Module):

    def __init__(self, input_dim, embedding_dim, hidden_dim, output_dim):

        super().__init__()
        self.embedding  = nn.Embedding(input_dim, embedding_dim)
        self.rnn        = nn.RNN(embedding_dim, hidden_dim)
        self.fc         = nn.Linear(hidden_dim, output_dim)

    def forward(self, text):

        embedding       = self.embedding(text)
        output, hidden  = self.rnn(embedding)

        return self.fc(hidden.squeeze(0))

如果您print()模型，这是输出。

RNN(
  (embedding): Embedding(25002, 100)
  (rnn): RNN(100, 256)
  (fc): Linear(in_features=256, out_features=1, bias=True)
)

以下是来自三种不同可视化工具的结果。

对于所有这些，您需要有可以通过模型的forward()方法的虚拟输入。 获取此输入的一种简单方法是从您的 Dataloader 中检索一个批次，如下所示：

batch = next(iter(dataloader_train))
yhat = model(batch.text) # Give dummy batch to forward().

火炬视点

我相信这个工具使用向后传递生成它的图，所以所有的框都使用 PyTorch 组件进行反向传播。

from torchviz import make_dot

make_dot(yhat, params=dict(list(model.named_parameters()))).render("rnn_torchviz", format="png")

此工具生成以下输出文件：

这是唯一明确提到我的模型中的三个层的输出， embedding 、 rnn和fc 。 运算符名称取自向后传递，因此其中一些难以理解。

隐藏层

我相信这个工具使用前向传递。

import hiddenlayer as hl

transforms = [ hl.transforms.Prune('Constant') ] # Removes Constant nodes from graph.

graph = hl.build_graph(model, batch.text, transforms=transforms)
graph.theme = hl.graph.THEMES['blue'].copy()
graph.save('rnn_hiddenlayer', format='png')

这是输出。 我喜欢蓝色的阴影。

我发现输出有太多细节并且混淆了我的架构。 例如，为什么unsqueeze被提到了这么多次？

耐创

此工具是适用于 Mac、Windows 和 Linux 的桌面应用程序。 它依赖于首先导出为的模型。 然后应用程序读取 ONNX 文件并呈现它。 然后有一个选项可以将模型导出到图像文件。

input_names = ['Sentence']
output_names = ['yhat']
torch.onnx.export(model, batch.text, 'rnn.onnx', input_names=input_names, output_names=output_names)

下面是模型在应用程序中的样子。 我认为这个工具非常灵活：您可以缩放和平移，还可以钻取图层和操作符。 我发现的唯一缺点是它只做垂直布局。