pytorch中model.save()保存的pth模型在Netron中不能画出结构图,因此转化成onnx模型再画
from torch.autograd import Variable
dummy_input = Variable(torch.randn(10, 3, 224, 224)).cuda()#randn(batch_size,channel,w,h)
torch.onnx.export(model, dummy_input, "model.onnx", export_params=True,opset_version=11)
如果模型中使用了upsample_bilinear2d上采样要加上opset_version=11,否则会报错:
ONNX export failed: Couldn’t export operator aten::upsample_bilinear2d
转化完成后使用Netron打开模型
参考:
https://zhuanlan.zhihu.com/p/197737707
https://pytorch.apachecn.org/docs/0.3/onnx.html
https://blog.csdn.net/info_black_hole/article/details/105556670