flytekitplugins.onnxpytorch.PyTorch2ONNXConfig

class flytekitplugins.onnxpytorch.PyTorch2ONNXConfig(args, export_params=True, verbose=False, training=<TrainingMode.EVAL: 0>, opset_version=9, input_names=<factory>, output_names=<factory>, operator_export_type=None, do_constant_folding=False, dynamic_axes=<factory>, keep_initializers_as_inputs=None, custom_opsets=<factory>, export_modules_as_functions=False)

PyTorch2ONNXConfig is the config used during the pytorch to ONNX conversion.

Parameters

• args (Union[Tuple, torch.Tensor]) – The input to the model.

• export_params (bool) – Whether to export all the parameters.

• verbose (bool) – Whether to print description of the ONNX model.

• training (torch._C._onnx.TrainingMode) – Whether to export the model in training mode or inference mode.

• opset_version (int) – The ONNX version to export the model to.

• input_names (List[str]) – Names to assign to the input nodes of the graph.

• output_names (List[str]) – Names to assign to the output nodes of the graph.

• operator_export_type (Optional[torch._C._onnx.OperatorExportTypes]) – How to export the ops.

• do_constant_folding (bool) – Whether to apply constant folding for optimization.

• dynamic_axes (Union[Dict[str, Dict[int, str]], Dict[str, List[int]]]) – Specify axes of tensors as dynamic.

• keep_initializers_as_inputs (Optional[bool]) – Whether to add the initializers as inputs to the graph.

• custom_opsets (Dict[str, int]) – A dictionary of opset doman name and version.

• export_modules_as_functions (Union[bool, set[Type]]) – Whether to export modules as functions.

Return type

None

Methods

classmethod from_dict(kvs, *, infer_missing=False)

Parameters

kvs (Optional[Union[dict, list, str, int, float, bool]]) –

Return type

dataclasses_json.api.A

classmethod from_json(s, *, parse_float=None, parse_int=None, parse_constant=None, infer_missing=False, **kw)

Parameters

s (Union[str, bytes, bytearray]) –

Return type

dataclasses_json.api.A

classmethod schema(*, infer_missing=False, only=None, exclude=(), many=False, context=None, load_only=(), dump_only=(), partial=False, unknown=None)

Parameters

• infer_missing (bool) –

• many (bool) –

• partial (bool) –

Return type

dataclasses_json.mm.SchemaF[dataclasses_json.mm.A]

to_dict(encode_json=False)

Return type

Dict[str, Optional[Union[dict, list, str, int, float, bool]]]

to_json(*, skipkeys=False, ensure_ascii=True, check_circular=True, allow_nan=True, indent=None, separators=None, default=None, sort_keys=False, **kw)

Parameters

• skipkeys (bool) –

• ensure_ascii (bool) –

• check_circular (bool) –

• allow_nan (bool) –

• indent (Optional[Union[int, str]]) –

• separators (Optional[Tuple[str, str]]) –

• default (Optional[Callable]) –

• sort_keys (bool) –

Return type

str

Attributes

do_constant_folding: bool = False

export_modules_as_functions: Union[bool, set[Type]] = False

export_params: bool = True

keep_initializers_as_inputs: Optional[bool] = None

operator_export_type: Optional[torch._C._onnx.OperatorExportTypes] = None

opset_version: int = 9

training: torch._C._onnx.TrainingMode = <TrainingMode.EVAL: 0>

verbose: bool = False

args: Union[Tuple, torch.Tensor]

input_names: List[str]

output_names: List[str]

dynamic_axes: Union[Dict[str, Dict[int, str]], Dict[str, List[int]]]

custom_opsets: Dict[str, int]