export¶

dragon.vm.torch.onnx.export(
  model,
  args,
  f,
  input_names=None,
  output_names=None,
  input_shapes=None,
  opset_version=None,
  verbose=False,
  enable_onnx_checker=True
)[source]¶

Export the recorded graph to an onnx model.

The outputs will be obtained by calling model(*args), both the tensor or numpy array are allowed:

class MyModule(torch.nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.fc = torch.nn.Linear(3, 3)

    def forward(self, x):
        y = self.fc(x)
        return y, np.ones((2, 3))

m = MyModule()
x = torch.zeros(2, 3)
torch.onnx.export(
    m,
    args=(x,),
    f='my_module.onnx',
    input_names=('x',),
    output_names=('y', 'ones'),
)

You can either specify the input_names, or pass a dict to the args. In the same way, model could return a dict to specify the output_names:

class MyModule(torch.nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.fc = torch.nn.Linear(3, 3)

    def forward(self, inputs):
        y = self.fc(inputs['x'])
        return {'y': y, 'ones': np.ones((2, 3))}

m = MyModule()
x = torch.zeros(2, 3)
torch.onnx.export(
    m,
    args={'x': x},
    f='my_module.onnx',
)

Also note that if a numpy array is given or returned, it’s name is definitely required. Otherwise, ONNX can’t export this value due to the lacking of id.

Parameters:

• model (dragon.vm.torch.nn.Module) – The module to export.
• args (Union[Sequence, Dict]) – The model inputs.
• f (str) – The filename for exporting model.
• input_names (Sequence[str], optional) – The name to the inputs.
• output_names (Sequence[str], optional) – The name to the outputs.
• input_shapes (Union[Sequence, Dict], optional) – The optional rewritten for input shapes.
• opset_version (int, optional) – The version of operator set.
• verbose (bool, optional, default=False) – Whether to print the debug string of graph.
• enable_onnx_checker (bool, optional, default=True) – Whether to check if model is valid.