如何在 Tensorboard 中可视化图神经网络的模型图

Posted 2023-03-27

【中文标题】如何在 Tensorboard 中可视化图神经网络的模型图

【英文标题】：How to Visualize the model graph of a Graph Neural Network in Tensorboard

【发布时间】：2021-10-24 11:25:16

【问题描述】：

我正在尝试可视化我为预测分子属性而制作的 Graph Neural Networks 的计算图。该模型是在 PyTorch 中制作的，并将 DGL 图作为输入。尝试可视化模型的代码 sn-p 如下所示：

train_log_dir = f'logs/datetime.datetime.now().strftime("%Y%m%d-%H%M%S")/train'
train_summary_writer = tensorboardX.SummaryWriter(train_log_dir)
train_summary_writer.add_graph(model, [transformer(dataset[0][0]), transformer(dataset[0][0])])

我遇到以下错误，TensorBoardX 无法可视化图模型，拒绝接受 DGL 图作为输入，只需要张量。有什么方法可以可视化模型吗？

RuntimeError: Tracer cannot infer type of (Graph(num_nodes=3, num_edges=4,
      ndata_schemes='x': Scheme(shape=(10,), dtype=torch.float32)
      edata_schemes='w': Scheme(shape=(4,), dtype=torch.float32)), Graph(num_nodes=3, num_edges=4,
      ndata_schemes='x': Scheme(shape=(10,), dtype=torch.float32)
      edata_schemes='w': Scheme(shape=(4,), dtype=torch.float32)))
:Only tensors and (possibly nested) tuples of tensors, lists, or dictsare supported as inputs or outputs of traced functions, but instead got value of type DGLHeteroGraph.

Process finished with exit code 1

【参考方案1】：

我通常使用 Torch 库中的 SummaryWriter。它的工作原理是这样的：

...
from torch.utils.tensorboard import SummaryWriter
...

# initializing your model

model = ...
dummy_input = ...

...
writer = SummaryWriter(f'logs/net')
writer.add_graph(model, dummy_input)

然后在终端运行你的python脚本后：

tensorboard --logdir logs

然后它会抛出类似 localhost:6006 的链接，然后就会出现您的可视化图形模型。 欲了解更多信息：https://pytorch.org/tutorials/intermediate/tensorboard_tutorial.html

【讨论】：

对，我就是这么做的，tensorboardX和torch tensorboard差不多，我也试过了。但是当我没有传递张量时问题就来了，而是我有图形对象。虚拟输入需要是张量或张量列表，以便 Torch 将其可视化。