This code is to export onnx file from computer vision pytorch model. Input type is related to single color image (Channel=3, Height, Width).
What shoud I change ?
① You should change model = net.Net() in get_model function. Net() means model's representative class name (ex. ResNet(), DenseNet() ...) in your model.py and Net() has necessary parameters you set.
② If you don't use input as single image, please revise dummy_data = torch.empty(1, 3, args.height, args.width) as your model's input type.
③ If your model has multiple outputs, please revise output_names of torch.onnx.export.
What does this code operate ?
① Load the deep learning network you want to use and set it to evaluation mode.
② Create the onnx model using the torch model.
③ Reload the created onnx model and compare the weights of the torch model and onnx model
④ If the onnx model has different weights from the existing torch model, update the whole and save it anew.
⑤ Finally, load the saved onnx model, add shape information, and then save it again.