构建torch 模型后,可视化计算流程可以复核自己是否写错,至关重要。
如下结构可视化方法,比summary()直观好用很多。
import netron modelpath = "./demo1.onnx" # 定义模型结构保存的路径 torch.onnx.export(self.model, input, modelpath) # 导出并保存 netron.start(modelpath)
但当你的环境在torch1.11.0 , torchvision 0.12.1以下,时候会提示,grid_sampler操作等不支持的问题。
!!!最直接的方法升级torch版本到1.12. 以上 (或者你再建一个环境里里面安装一个适应的版本),如果还提示,请如下图修改
opset_version的值
真的比其他方法要一步到位,且不会出现其他问题。