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

Question

我正在尝试可视化我用来预测分子属性的图神经网络的计算图。该模型是在 PyTorch 中制作的，并采用 DGL 图作为输入。尝试可视化模型的代码片段如下所示：

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

Answer 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

Answer 2

你现在解决这个问题了吗？或者您是否找到任何其他工具可以在 DGL 中可视化 gnn 模型？