medusa.onnx#

Module with a class to make working with ONNX models easier.

Module Contents#

class medusa.onnx.OnnxModel(onnx_file, device=DEVICE, **kwargs)[source]#

Wrapper around an onnxruntime session to make running with an iobinding easier.

Parameters:

  • onnx_file (str, pathlib.Path) – Path to onnx file

  • device (str) – Device to run model on (‘cpu’ or ‘cuda’)

  • **kwargs – Extra keyword arguments (from ‘in_names’, ‘in_shapes’, ‘out_names’, ‘out_shapes’) to be set in _params (which override those parameters from the onnx model)

run(inputs, outputs_as_list=False)[source]#

Runs the model with given inputs.

Parameters:

  • inputs (torch.tensor, list) – Either a list of torch tensors or a single tensor (if just one input)

  • outputs_as_list (bool) – Whether to return a list of outputs instead of a dict

Returns:

outputs – Dictionary (or, when outputs_as_list, a list) with outputs names and corresponding outputs (as torch tensors)

Return type:

OrderedDict, list