def _jit_graph_to_onnx_model(graph, operator_export_type, opset_version): r""" This function exports torch::jit::Graph object to serialized ONNX ModelProto. This function is for testing purpose. It only keeps the essential parts for IR graph conversions. It also does not interact with actual PyTorch modules nor PyTorch tensor inputs. """ from torch.onnx.symbolic_helper import _set_onnx_shape_inference, _set_opset_version from torch.onnx.utils import _optimize_graph # Shape inference is required because some ops' symbolic functions # generate sub-graphs based on inputs' types. _set_onnx_shape_inference(True) _set_opset_version(opset_version) graph = _optimize_graph(graph, operator_export_type, params_dict={}) proto, _, _, _ = graph._export_onnx( {}, opset_version, {}, False, operator_export_type, False, False, {}, True, "", {}, ) return proto
def _model_to_graph(self, model, input, do_constant_folding=True, training=TrainingMode.EVAL, operator_export_type=OperatorExportTypes.ONNX, input_names=None, dynamic_axes=None): if training == torch.onnx.TrainingMode.TRAINING: model.train() elif training == torch.onnx.TrainingMode.EVAL: model.eval() # Need disable onnx_shape_inference for this test because it puts const node to initializers. _set_onnx_shape_inference(False) utils._validate_dynamic_axes(dynamic_axes, model, None, None) graph, params_dict, torch_out = utils._model_to_graph( model, input, do_constant_folding=do_constant_folding, _disable_torch_constant_prop=True, operator_export_type=operator_export_type, training=training, input_names=input_names, dynamic_axes=dynamic_axes) _set_onnx_shape_inference(True) return graph, params_dict, torch_out
def _convert(self) -> None: prev_opset_version = None prev_export_type = None prev_shape_inference = None try: assert not to_utils.is_in_onnx_export() # type: ignore[no-untyped-call] with to_utils.select_model_mode_for_export(self.original_model, self.training): to_utils.__IN_ONNX_EXPORT = True prev_opset_version = sym_hel._export_onnx_opset_version sym_hel._set_opset_version(self.opset_version) # type: ignore[no-untyped-call] prev_export_type = sym_hel._operator_export_type sym_hel._set_operator_export_type(self.operator_export_type) # type: ignore[no-untyped-call] prev_shape_inference = sym_hel._onnx_shape_inference sym_hel._set_onnx_shape_inference( # type: ignore[no-untyped-call] False # TODO(twata): Use `self.onnx_shape_inference` ) self._run_trace() self.model: onnx.ModelProto = self.generate_onnx() finally: to_utils.__IN_ONNX_EXPORT = False if prev_opset_version is not None: sym_hel._set_opset_version(prev_opset_version) # type: ignore[no-untyped-call] if prev_shape_inference is not None: sym_hel._set_operator_export_type(prev_export_type) # type: ignore[no-untyped-call] if prev_shape_inference is not None: sym_hel._set_onnx_shape_inference(prev_shape_inference) # type: ignore[no-untyped-call]
def run_model_test(self, model, batch_size=2, state_dict=None, input=None, use_gpu=True, rtol=0.001, atol=1e-7, example_outputs=None, do_constant_folding=True, dynamic_axes=None, test_with_inputs=None, input_names=None, output_names=None, fixed_batch_size=False): model.eval() if input is None: input = torch.randn(batch_size, 3, 224, 224, requires_grad=True) with torch.no_grad(): if isinstance(input, torch.Tensor): input = (input, ) # In-place operators will update input tensor data as well. # Thus inputs are replicated before every forward call. input_copy = copy.deepcopy(input) output = model(*input_copy) if isinstance(output, torch.Tensor): output = (output, ) _set_opset_version(self.opset_version) _set_operator_export_type(OperatorExportTypes.ONNX) _set_onnx_shape_inference(True) _set_training_mode(False) if dynamic_axes is None: dynamic_axes = {} _validate_dynamic_axes(dynamic_axes, model, input_names, output_names) input_copy = copy.deepcopy(input) graph, _, _ = utils._model_to_graph( model, input_copy, input_names=input_names, output_names=output_names, operator_export_type=OperatorExportTypes.ONNX, example_outputs=output, do_constant_folding=do_constant_folding, training=TrainingMode.EVAL, use_new_jit_passes=self.use_new_jit_passes, dynamic_axes=dynamic_axes) verify_inferred_shape(graph)
def _export_jit_graph_to_onnx_model_proto(graph: torch._C.Graph, operator_export_type: int): from torch.onnx.symbolic_helper import _set_onnx_shape_inference, _set_operator_export_type, _set_opset_version _set_onnx_shape_inference(True) _set_operator_export_type(operator_export_type) torch._C._jit_pass_run_decompositions(graph) graph = torch.onnx.utils._optimize_graph(graph, operator_export_type, params_dict={}) proto, _, _, _ = graph._export_onnx( {}, torch.onnx._globals.GLOBALS.export_onnx_opset_version, {}, False, operator_export_type, False, False, {}, True, "", {}, ) return proto
def __init__(self, *args, **kwargs): unittest.TestCase.__init__(self, *args, **kwargs) self.opset_version = _onnx_main_opset _set_onnx_shape_inference(True) _set_opset_version(self.opset_version)
def setUp(self): self.opset_version = _constants.onnx_main_opset symbolic_helper._set_onnx_shape_inference(True) symbolic_helper._set_opset_version(self.opset_version)
def __init__(self, *args, **kwargs): TestCase.__init__(self, *args, **kwargs) self.opset_version = _constants.onnx_main_opset _set_onnx_shape_inference(True) _set_opset_version(self.opset_version)