tensorboard踩坑指南add_graph

想用tensorboard把网络可视化一下，因为我的网络有多个通道的输入，因此需要用括号括起来

from torchinfo import summary
summary(model, [(10, 1, 16, 16), (10, 1)])
dummy_input01 = torch.rand(10, 1, 16, 16).to(device)  #假设输入10张1*16*16的图片

dummy_input02 = torch.rand(10, 1).to(device)

from torch.utils.tensorboard import SummaryWriter
with SummaryWriter(comment='Net') as w:
    w.add_graph(model, (dummy_input01, dummy_input02))

但是这个时候报错了，summary原来是直接用的torchsummary库，好像不支持维度为1的数据，因此换成了torchinfo库，就好了。

还有一个报错“ NYI: Named tensors are not supported with the tracer ”，在跟踪器中不支持命名张量。这个问题解决了好久，把tensor进行unbind、去掉梯度、rename(None)都没有用。

后来查到，只需要将模型转换为脚本模式，以避免跟踪器的限制。也就是在写进log之前加一行：

model = torch.jit.script(model)

问题完美解决！