def test_constant_fold_concat(self): class ConcatModule(torch.nn.Module): def forward(self, x): # Why did I insert a Cast here? There appears to be intentional # behavior in ONNX constant folding where constant tensors which # are not attached to any known to be foldable onnx # operations don't get extracted into the initializer graph. So # without these casts, we will actually fail to pull out one of # the constants, thus failing constant folding. I think the # test is wrong but I don't have time to write a more correct # test (I think the right way to go about the test is to setup # a predicate for what invariant graphs should hold after # constant folding, and then verify this predicate holds. # I think the asserts below are an attempt at this predicate, # but it is not right!) # # More commentary at # https://github.com/pytorch/pytorch/pull/18698/files#r340107552 a = torch.tensor([[1., 2., 3.]]).to(torch.float) b = torch.tensor([[4., 5., 6.]]).to(torch.float) c = torch.cat((a, b), 0) d = b + c return x + d _set_opset_version(self.opset_version) _set_operator_export_type(OperatorExportTypes.ONNX) x = torch.ones(2, 3) graph, _, __ = utils._model_to_graph( ConcatModule(), (x, ), do_constant_folding=True, _disable_torch_constant_prop=True, operator_export_type=OperatorExportTypes.ONNX) for node in graph.nodes(): assert node.kind() != "onnx::Concat" assert node.kind() != "onnx::Cast" assert node.kind() != "onnx::Constant" assert len(list(graph.nodes())) == 1
def test_constant_fold_sub(self): class Module(torch.nn.Module): def __init__(self, ): super(Module, self).__init__() self.register_buffer("weight", torch.ones(5)) def forward(self, x): sub = self.weight - torch.tensor([1, 2, 3, 4, 5]) return sub + x x = torch.randn(2, 5) _set_opset_version(self.opset_version) _set_operator_export_type(OperatorExportTypes.ONNX) graph, params_dict, __ = utils._model_to_graph( Module(), (x, ), do_constant_folding=True, operator_export_type=OperatorExportTypes.ONNX) for node in graph.nodes(): assert node.kind() != "onnx::Sub" self.assertEqual(len(list(graph.nodes())), 1) params = list(params_dict.values()) self.assertEqual(len(params), 1) weight = params[0] self.assertEqual(weight, torch.tensor([0, -1, -2, -3, -4]))
def test_unused_initializers(self): class Model(torch.nn.Module): def __init__(self): super(Model, self).__init__() self.conv2 = torch.nn.ConvTranspose2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2), dilation=(1, 1)) self.k_proj = torch.nn.Linear(5, 5, bias=True) def forward(self, x): x = self.conv2(x) return x x = torch.randn(20, 16, 50, 100) _set_opset_version(self.opset_version) _set_operator_export_type(OperatorExportTypes.ONNX) _, params_dict, __ = utils._model_to_graph( Model(), (x, ), do_constant_folding=False, operator_export_type=OperatorExportTypes.ONNX) assert len(params_dict) == 2
def test_constant_fold_add(self): class Module(torch.nn.Module): def __init__(self, ): super(Module, self).__init__() self.register_buffer("weight", torch.ones(5)) def forward(self, x): add = self.weight + torch.tensor([1, 2, 3, 4, 5]) return add - x x = torch.randn(2, 5) _set_opset_version(self.opset_version) _set_operator_export_type(OperatorExportTypes.ONNX) graph, params_dict, __ = utils._model_to_graph( Module(), (x, ), do_constant_folding=True, operator_export_type=OperatorExportTypes.ONNX) for node in graph.nodes(): self.assertTrue(node.kind() != "onnx::Add") self.assertEqual(len(list(graph.nodes())), 1) params = list(params_dict.values()) self.assertEqual(len(params), 1) weight = params[0] # TODO(#38095): Replace assertEqualIgnoreType. See issue #38095 self.assertEqualIgnoreType(weight, torch.tensor([2, 3, 4, 5, 6]))
def test_scripting_param(self): class MyModule(torch.nn.Module): def __init__(self): super(MyModule, self).__init__() self.conv = torch.nn.Conv2d(3, 16, kernel_size=1, stride=2, padding=3, bias=True) self.bn = torch.nn.BatchNorm2d(16, affine=True) def forward(self, x): x = self.conv(x) bn = self.bn(x) return bn model = torch.jit.script(MyModule()) x = torch.randn(10, 3, 128, 128) example_outputs = model(x) f = io.BytesIO() _set_opset_version(self.opset_version) _set_operator_export_type(OperatorExportTypes.ONNX) graph, _, __ = utils._model_to_graph(model, (x,), do_constant_folding=True, example_outputs=example_outputs, operator_export_type=OperatorExportTypes.ONNX) graph_input_params = [param.debugName() for param in graph.inputs()] assert all(item in graph_input_params for item in dict(model.named_parameters())), \ "Graph parameter names does not match model parameters."
def _model_to_graph(self, model, input, do_constant_folding=True, example_outputs=None, 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, example_outputs=example_outputs, input_names=input_names, dynamic_axes=dynamic_axes) _set_onnx_shape_inference(True) return graph, params_dict, torch_out
def _onnx_graph_from_model( model: Union[torch.nn.Module, torch.jit.ScriptModule], args: Tuple[Any, ...], kwargs: Mapping[str, Any], export_options: _experimental.ExportOptions, ) -> _C.Graph: """As part of the ONNX export steps, export an ONNX JIT graph from a PyTorch model. Args: model: See :func:`check_export_model_diff`. args: See :func:`check_export_model_diff`. kwargs: See :func:`check_export_model_diff`. export_options: See :func:`check_export_model_diff`. Returns: onnx_graph (_C.Graph): An ONNX JIT graph. """ # TODO: refactor utils.py to remove duplicated code of context setup. See #78834 opset_version = export_options.opset_version operator_export_type = export_options.operator_export_type export_modules_as_functions = export_options.export_modules_as_functions training = export_options.training verbose = export_options.verbose dynamic_axes = export_options.dynamic_axes input_names = export_options.input_names output_names = export_options.output_names if opset_version is None: opset_version = _constants.onnx_default_opset export_modules_as_functions = utils._setup_trace_module_map( model, export_modules_as_functions ) if not operator_export_type: if _C_onnx._CAFFE2_ATEN_FALLBACK: operator_export_type = _C_onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK else: operator_export_type = _C_onnx.OperatorExportTypes.ONNX GLOBALS.export_onnx_opset_version = opset_version GLOBALS.operator_export_type = operator_export_type with utils.exporter_context(model, training, verbose): do_constant_folding = utils._decide_constant_folding( export_options.do_constant_folding, operator_export_type, training ) if dynamic_axes is None: dynamic_axes = {} utils._validate_dynamic_axes(dynamic_axes, model, input_names, output_names) export_inputs = _prepare_input_for_export(args, kwargs) export_inputs = utils._decide_input_format(model, export_inputs) onnx_graph, _, _ = utils._model_to_graph( model, export_inputs, verbose, input_names, output_names, operator_export_type, do_constant_folding, training=training, dynamic_axes=dynamic_axes, ) return onnx_graph
def pytorch_to_mdf( model: Union[Callable, torch.nn.Module, torch.ScriptFunction, torch.ScriptModule], args: Union[None, torch.Tensor, Tuple[torch.Tensor]] = None, example_outputs: Union[None, torch.Tensor, Tuple[torch.Tensor]] = None, trace: bool = False, use_onnx_ops: bool = True, ) -> Union[Model, Graph]: r""" Convert a PyTorch model to an MDF model. By default, this function will invoke `torch.jit.script` on the model to compile it down to TorchScript IR and simplify the graph before exporting the MDF. The default is to use ONNX operations when possible and fallback to ATEN\Torch ops when ONNX support is not available (`torch._C._onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK` mode). To use allATEN\Torch ops, set use_onnx_ops to False. Args: model: The model to translate into MDF. args: The input arguments for this model. If a nn.Module is passed then the model will be traced with these inputs. If a ScriptModule is passed, they are still needed to deterimine input shapes. example_outputs: Example outputs from the model for determing output shapes. trace: Force the use of tracing to compile the model. The default is to use torch.jit.script use_onnx_ops: Use ONNX ops when possible, fallback to ATEN ops when not available. Default is True. If False, use only ATEN ops. Returns: The translated MDF model """ # Get the graph and nodes from the TorchScript model try: # If the graph attribute is available, we are dealing with a already jitted model (ScriptModule, ScriptFunciton, # etc.) graph = model.graph jit_model = model except AttributeError: # Lets jit things, if the user doesn't want to trace or we are dealing with a standard Python function, we need # to JIT script it. if not trace or inspect.isfunction(model): jit_model = torch.jit.script(model) graph = jit_model.graph else: # If the user wants to trace, _model_to_graph below will take care of that for us. graph = None if use_onnx_ops: operator_export_type = torch._C._onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK else: operator_export_type = torch._C._onnx.OperatorExportTypes.RAW # Call out to a part of the ONNX exporter that simiplifies the graph before ONNX export. from torch.onnx.utils import _model_to_graph from torch.onnx import TrainingMode from torch.onnx.symbolic_helper import ( _export_onnx_opset_version, _set_opset_version, ) previous_opset_version = _export_onnx_opset_version _set_opset_version(modeci_onnx_opset_version) graph, params_dict, torch_out = _model_to_graph( model=jit_model if graph else model, args=args, example_outputs=example_outputs, do_constant_folding=False, training=TrainingMode.EVAL, _retain_param_name=True, operator_export_type=operator_export_type, dynamic_axes={}, ) _set_opset_version(previous_opset_version) model_name, graph_name = make_model_graph_name(model) # Setup the MDF model and graph mdf_model = Model(id=model_name) mdf_graph = Graph(id=graph_name) mdf_model.graphs.append(mdf_graph) # Get all constant nodes in the graph consts = get_graph_constants(graph) # Get any inputs to the graph, and their debug names. Pass args so we know how # many original input arguments the graph has. ONNX lowering from _model_to_graph # makes all parameters to the model inputs. port_mapper = PortMapper(graph=graph, args=args) # Translate the TorchScript graph to and MDF graph object. This could be a recursive call translate_graph(graph=graph, mdf_graph=mdf_graph, consts=consts, port_mapper=port_mapper) # Replace in "." for "_" in parameter names. We have done this elsewhere when creating the input ports for these # parameters. params_dict = { port_mapper.id_to_port(k): v for k, v in params_dict.items() } # Set the ONNX opset version mdf_model.onnx_opset_version = _export_onnx_opset_version return mdf_model, params_dict