Convert custom model to ONNX format

Pytorch

Steps based on EXPORTING A MODEL FROM PYTORCH TO ONNX AND RUNNING IT USING ONNX RUNTIME.

• Step 0. Requirements:

  • Pytorch

  • ONNX

• Step 1. Load PyTorch model .. code-block:

  from torch import nn
  import torch.utils.model_zoo as model_zoo
  import torch.onnx
  
  model = ... # your model instation
  model.load_state_dict(torch.load(YOUR_MODEL_CHECKPOINT_PATH, map_location='cpu')['state_dict'])
  model.eval()
  

• Step 2. Create data sample with batch_size=1 and call forward step of your model: .. code-block:

  x = torch.rand(1, INP_CHANNEL, INP_HEIGHT, INP_WIDTH) # eg. torch.rand([1, 3, 256, 256])
  _ = model(x)
  

• Step 3. Call export function

  torch.onnx.export(model,
                  x,  # model input
                  'model.onnx',  # where to save the model
                  export_params=True,
                  opset_version=15,
                  input_names=['input'],
                  output_names=['output'],
                  do_constant_folding=False)
  

Tensorflow/Keras

Steps based on the tensorflow-onnx repository. The instruction is valid for saved model format. For other types follow tensorflow-onnx instructions.

• Requirements:

  • tensorflow

  • ONNX

  • tf2onnx

• And simply call converter script:

  python -m tf2onnx.convert --saved-model YOUR_MODEL_CHECKPOINT_PATH --output model.onnx --opset 15