pytorch模型转换为onnx，并使用netron可视化

netron 是一个非常好用的网络结构可视化工具。

但是netron对pytorch模型的支持还不成熟。自己试的效果是生成的模型图没有连线。

目前支持的框架 根据netron的github

目前netron支持：

ONNX (.onnx, .pb, .pbtxt)
Keras (.h5, .keras)
Core ML (.mlmodel)
Caffe (.caffemodel, .prototxt)
Caffe2 (predict_net.pb, predict_net.pbtxt)
Darknet (.cfg)
MXNet (.model, -symbol.json)
ncnn (.param) 
TensorFlow Lite (.tflite)
PaddlePaddle (.zip, model)
TensorFlow.js
CNTK (.model, .cntk)

并且实验性支持:

TorchScript (.pt, .pth)
PyTorch (.pt, .pth)
Torch (.t7)
Arm NN (.armnn)
BigDL (.bigdl, .model) 
Chainer (.npz, .h5)
Deeplearning4j (.zip)
MediaPipe (.pbtxt)
ML.NET (.zip), MNN (.mnn)
OpenVINO (.xml)
scikit-learn (.pkl)
TensorFlow (.pb, .meta, .pbtxt, .ckpt, .index)

Netron supports ONNX, TensorFlow Lite, Caffe, Keras, Darknet, PaddlePaddle, ncnn, MNN, Core ML, RKNN, MXNet, MindSpore Lite, TNN, Barracuda, Tengine, CNTK, TensorFlow.js, Caffe2 and UFF.

Netron has experimental support for PyTorch, TensorFlow, TorchScript, OpenVINO, Torch, Vitis AI, kmodel, Arm NN, BigDL, Chainer, Deeplearning4j, MediaPipe, ML.NET and scikit-learn.

这里就有一个把 .pth 模型转化为 .onnx 模型。

Pytorch模型转onnx

model = resnet18(pretrained=True)
# print(model)
# old_net_path = "resnet18.pth"
new_net_path = "./resnet18.onnx"

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 导入模型
net = model.to(device)
# net.load_state_dict(torch.load(old_net_path, map_location=device))
net.eval()

input = torch.randn(1, 3, 224, 224).to(device)  # BCHW  其中Batch必须为1，因为测试时一般为1，尺寸HW必须和训练时的尺寸一致
torch.onnx.export(net, input, new_net_path, verbose=False)

标签：netron,onnx,torch,pytorch,TensorFlow,model,net
来源： https://www.cnblogs.com/michaelcjl/p/16576323.html