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:
export_params (bool) – Whether to export all the parameters.
verbose (bool) – Whether to print description of the ONNX model.
training (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 (OperatorExportTypes | None) – How to export the ops.
do_constant_folding (bool) – Whether to apply constant folding for optimization.
dynamic_axes (Dict[str, Dict[int, str]] | Dict[str, List[int]]) – Specify axes of tensors as dynamic.
keep_initializers_as_inputs (bool | None) – Whether to add the initializers as inputs to the graph.
custom_opsets (Dict[str, int]) – A dictionary of opset domain name and version.
export_modules_as_functions (bool | set[Type]) – Whether to export modules as functions.
Methods
- classmethod from_dict(kvs, *, infer_missing=False)#
- classmethod from_json(s, *, parse_float=None, parse_int=None, parse_constant=None, infer_missing=False, **kw)#
- classmethod schema(*, infer_missing=False, only=None, exclude=(), many=False, context=None, load_only=(), dump_only=(), partial=False, unknown=None)#
- to_json(*, skipkeys=False, ensure_ascii=True, check_circular=True, allow_nan=True, indent=None, separators=None, default=None, sort_keys=False, **kw)#
Attributes
- dataclass_json_config = None
- do_constant_folding: bool = False
- export_params: bool = True
- operator_export_type: OperatorExportTypes | None = None
- opset_version: int = 9
- training: TrainingMode = <TrainingMode.EVAL: 0>
- verbose: bool = False