pytorch - 使用 torchviz make_dots 显示具有多个输出的 PyTorch 模型

Question

我有一个具有多个输出的模型，确切地说是 4 个：

 def forward(self, x):
      outputs = []
      for conv, act in zip(self.Convolutions, self.Activations):
           y = conv(x)
           outputs.append(act(y))
      return outputs

我想使用make_dotfrom显示它torchviz：

 from torchviz import make_dot
 generator = ...
 batch = next(iter(generator))
 input, output = batch["input"].to(device, dtype=torch.float), batch["output"].to(device, dtype=torch.float)
 dot = make_dot(model(input), params=dict(model.named_parameters()))

但我收到以下错误：

 File "/opt/local/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/torchviz/dot.py", line 37, in make_dot
 output_nodes = (var.grad_fn,) if not isinstance(var, tuple) else tuple(v.grad_fn for v in var)
 AttributeError: 'list' object has no attribute 'grad_fn'

显然列表没有grad_fn功能，但根据这个讨论，我可以返回一个输出列表。

我究竟做错了什么？

Answer 1

模型可以返回一个列表，但make_dot想要一个Tensor. 如果输出组件具有相似的形状，我建议使用torch.cat它。