def test_aten_embedding_1(self): _onnx_opset_version = 12 @parse_args('v', 'v', 'i', 'b', 'b') def embedding(g, weight, indices, padding_idx, scale_grad_by_freq, sparse): custom_attributes_json = ( '{' f'"padding_idx":{str(padding_idx)},' f'"scale_grad_by_freq":{str(scale_grad_by_freq).lower()},' f'"sparse":{str(sparse).lower()}' '}' ) output = g.op("com.microsoft::ATenOp", weight, indices, name_s='aten::embedding', custom_attributes_json_s=custom_attributes_json) return output register_custom_op_symbolic('::embedding', embedding, _onnx_opset_version) class Model(torch.nn.Module): def __init__(self): super().__init__() self.emb = torch.nn.Embedding(4, 8) def forward(self, x, y): res = self.emb(x) res = res + y return torch.ones(res.shape[0]) model = Model() x = torch.ones(32, dtype=torch.long) y = torch.randn(1, 8) self.assertONNX(model, (x, y), opset_version=_onnx_opset_version) unregister_custom_op_symbolic('::embedding', _onnx_opset_version)
def test_aten_embedding_2(self): _onnx_opset_version = 12 @parse_args('v', 'v', 'i', 'b', 'b') def embedding(g, weight, indices, padding_idx, scale_grad_by_freq, sparse): custom_attributes_json = ( '{' f'"padding_idx":{str(padding_idx)},' f'"scale_grad_by_freq":{str(scale_grad_by_freq).lower()},' f'"sparse":{str(sparse).lower()}' '}') output = g.at("embedding", weight, indices, custom_attributes_json_s=custom_attributes_json) # do shape inference and set it via setType indices_shape = _get_tensor_sizes(indices) if indices_shape is not None and hasattr(weight.type(), 'with_sizes'): output_type = weight.type().with_sizes( indices_shape + [_get_tensor_dim_size(weight, 1)]) output.setType(output_type) return output register_custom_op_symbolic('::embedding', embedding, _onnx_opset_version) class Model(torch.nn.Module): def __init__(self): super().__init__() self.emb = torch.nn.Embedding(4, 8) def forward(self, x, y): res = self.emb(x) res = res + y return torch.ones(res.shape[0]) model = Model() x = torch.ones(32, dtype=torch.long) y = torch.randn(1, 8) self.assertONNX(model, (x, y), opset_version=_onnx_opset_version, input_names=['input_1', 'input_2'], dynamic_axes={ "input_1": { 0: "dim_0" }, 'input_2': { 0: "dim_1", 1: "dim_2" } }, keep_initializers_as_inputs=False, operator_export_type=torch.onnx.OperatorExportTypes. ONNX_ATEN_FALLBACK) unregister_custom_op_symbolic('::embedding', _onnx_opset_version)
def test_aten_embedding_2(self): _onnx_opset_version = 12 @parse_args("v", "v", "i", "b", "b") def embedding(g, weight, indices, padding_idx, scale_grad_by_freq, sparse): custom_attributes_json = ( "{" f'"padding_idx":{str(padding_idx)},' f'"scale_grad_by_freq":{str(scale_grad_by_freq).lower()},' f'"sparse":{str(sparse).lower()}' "}") output = g.op("com.microsoft::ATenOp", weight, indices, name_s="aten::embedding", custom_attributes_json_s=custom_attributes_json) # do shape inference and set it via setType indices_shape = _get_tensor_sizes(indices) if indices_shape is not None and hasattr(weight.type(), "with_sizes"): output_type = weight.type().with_sizes( indices_shape + [_get_tensor_dim_size(weight, 1)]) output.setType(output_type) return output register_custom_op_symbolic("::embedding", embedding, _onnx_opset_version) class Model(torch.nn.Module): def __init__(self): super().__init__() self.emb = torch.nn.Embedding(4, 8) def forward(self, x, y): res = self.emb(x) res = res + y return torch.ones(res.shape[0]) model = Model() x = torch.ones(32, dtype=torch.long) y = torch.randn(1, 8) self.assertONNX(model, (x, y), opset_version=_onnx_opset_version, input_names=["input_1", "input_2"], dynamic_axes={ "input_1": { 0: "dim_0" }, "input_2": { 0: "dim_1", 1: "dim_2" } }) unregister_custom_op_symbolic("::embedding", _onnx_opset_version)
def enable_custom_autograd_support(enable=True): import atexit from torch.onnx import register_custom_op_symbolic, unregister_custom_op_symbolic from onnxruntime.capi._pybind_state import ( register_backward_runner, register_forward_runner, unregister_python_functions, ) from onnxruntime.training.ortmodule.torch_cpp_extensions import torch_interop_utils from ._custom_autograd_function_exporter import _export from ._custom_autograd_function_runner import call_python_backward_function, call_python_forward_function if enable is True: if custom_autograd_function_enabler.already_enabled is False: # Initialize static objects needed to run custom autograd.Function's. register_forward_runner(call_python_forward_function) register_backward_runner(call_python_backward_function) # Unregister all python functions automatically upon normal interpreter termination. atexit.register(unregister_python_functions) # Clear all gradient functions, to avoid a deadlock issue. # Check the called function for more detailed comments. atexit.register(torch_interop_utils.clear_all_grad_fns) try: # This is for the latest Pytorch nightly after this commit: # https://github.com/pytorch/pytorch/commit/11bc435622e6b7207bbf37ed1aafe999e1f296ec register_custom_op_symbolic("prim::PythonOp", _export, 1) except: # This applies to Pytorch 1.9 and 1.9.1. register_custom_op_symbolic("::prim_PythonOp", _export, 1) custom_autograd_function_enabler.state = True else: if custom_autograd_function_enabler.already_enabled is True: # We don't need remove the registered runner because it won't be used if we disable the feature. # But we need unregister the PythonOp custom operator function. try: # This is for the latest Pytorch nightly after this commit: # https://github.com/pytorch/pytorch/commit/11bc435622e6b7207bbf37ed1aafe999e1f296ec unregister_custom_op_symbolic("prim::PythonOp", 1) except: # This applies to Pytorch 1.9 and 1.9.1. unregister_custom_op_symbolic("::prim_PythonOp", 1) custom_autograd_function_enabler.state = False
def test_aten_embedding_1(self): _onnx_opset_version = 12 @parse_args("v", "v", "i", "b", "b") def embedding(g, weight, indices, padding_idx, scale_grad_by_freq, sparse): custom_attributes_json = ( "{" f'"padding_idx":{str(padding_idx)},' f'"scale_grad_by_freq":{str(scale_grad_by_freq).lower()},' f'"sparse":{str(sparse).lower()}' "}") output = g.at( "embedding", weight, indices, custom_attributes_json_s=custom_attributes_json, ) return output register_custom_op_symbolic("::embedding", embedding, _onnx_opset_version) class Model(torch.nn.Module): def __init__(self): super().__init__() self.emb = torch.nn.Embedding(4, 8) def forward(self, x, y): res = self.emb(x) res = res + y return torch.ones(res.shape[0]) model = Model() x = torch.ones(32, dtype=torch.long) y = torch.randn(1, 8) self.assertONNX(model, (x, y), opset_version=_onnx_opset_version) unregister_custom_op_symbolic("::embedding", _onnx_opset_version)