module tools.onnx_helper#

Short summary#

module deeponnxcustom.tools.onnx_helper

Helpers about ONNX.

source on GitHub

Functions#

function

truncated documentation

onnx_rename_weights

Renames ONNX initialiers to make sure their name follows the alphabetical order. The model is modified inplace. …

save_as_onnx

Converts a torch model into ONNX using torch.onnx.export(). The function works on models with only one input. …

Documentation#

Helpers about ONNX.

source on GitHub

deeponnxcustom.tools.onnx_helper.onnx_rename_weights(onx)#

Renames ONNX initialiers to make sure their name follows the alphabetical order. The model is modified inplace. This function calls onnx_rename_names.

Parameters

onx – ONNX model

Returns

same model

Note

The function does not go into subgraphs.

source on GitHub

deeponnxcustom.tools.onnx_helper.save_as_onnx(model, filename, size=None, target_opset=14, batch_size=1, device='cpu', keep_initializers_as_inputs=False)#

Converts a torch model into ONNX using torch.onnx.export(). The function works on models with only one input.

Parameters
  • model – torch model

  • filename – output filename

  • size – input size or left None to guess it from the model

  • target_opset – opset to use for the conversion

  • batch_size – batch size

  • device – device

  • keep_initializers_as_inputs – see torch.onnx.export()

Export a torch model into ONNX

import torch
from deeponnxcustom.tools.onnx_helper import save_as_onnx

class MyModel(torch.nn.Module):
    # ...

nn = MyModel()
save_as_onnx(nn, "my_model.onnx")

source on GitHub