module tools.onnx_helper
#
Short summary#
module deeponnxcustom.tools.onnx_helper
Helpers about ONNX.
Functions#
function |
truncated documentation |
---|---|
Renames ONNX initialiers to make sure their name follows the alphabetical order. The model is modified inplace. … |
|
Converts a torch model into ONNX using |
Documentation#
Helpers about ONNX.
- 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.
- 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")