不同于tensorflow提供了强大的可视化功能,pytorch可视化一般借助第三方工具,更通用的方法是使用tensorboardX。
https://github.com/lanpa/tensorboardX
使用命令:pip install tensorboardX
安装后,如下代码可视化模型:
print(model)
writer = SummaryWriter(log_dir="./logs/",comment="myresnet")
with writer:
writer.add_graph(model,input_to_model=torch.rand(1,3,32,32))
更多内容,访问文档:
https://tensorboardx.readthedocs.io/en/latest/index.html