convert_to_onnx.py

import numpy as np
import cv2
import torch
import onnx
import os

if __name__ == '__main__':
    torch.set_grad_enabled(False)
    checkpoint = torch.load("path_to_saved_model")
    import torch.nn as nn

    model = checkpoint['model'].module
    model.eval()
    device = torch.device('cpu')
    model = model.to(device)
    model = model.eval()
    # print("Load model : Done ")
    # ##################export###############
    output_onnx = "model.onnx"
    print("==> Exporting model to ONNX format at '{}'".format(output_onnx))
    input_names = ["input0"]
    output_names = ["output0"]
    inputs = torch.randn(1, 3, 112, 112)
    torch_out = torch.onnx._export(model,
                                   inputs,
                                   output_onnx,
                                   verbose=True,
                                   input_names=input_names,
                                   output_names=output_names,
                                   example_outputs=True,  # to show sample output dimension
                                   keep_initializers_as_inputs=True,  # to avoid error _Map_base::at
                                   # opset_version=7, # need to change to 11, to deal with tensorflow fix_size input
                                   # dynamic_axes={
                                   #     "input0": [2, 3],
                                   #     "loc0": [1, 2],
                                   #     "conf0": [1, 2],
                                   #     "landmark0": [1, 2]
                                   # }
                                   )
    # Simplify model
    print("Simplifying onnx model!")
    onnx_simpli = output_onnx[:-5] + "_simpli.onnx"
    simpli_cmd = "python -m onnxsim --skip-fuse-bn " + output_onnx + " " + onnx_simpli
    os.system(simpli_cmd)