def test_contrib_op_with_loop(self): class M(torch.nn.Module): def __init__(self): super().__init__() self.gelu = torch.nn.GELU(approximate="none") def forward(self, x): res = [] res2 = [] for i in range(x.size(0)): if len(res) > 0: res2.append(res[0]) else: res2.append(self.gelu(x[0])) res.append(x[0]) return torch.stack(res), torch.stack(res2) def symbolic_custom_gelu(g, input, approximate): return g.op("com.microsoft::Gelu", input).setType(input.type()) from torch.onnx import register_custom_op_symbolic register_custom_op_symbolic("::gelu", symbolic_custom_gelu, 1) x = torch.randn(3, 3, 4, requires_grad=True) model = torch.jit.script(M()) onnx_test_common.run_model_test(self, model, input_args=(x, ))
def exportTest( self, model, inputs, rtol=1e-2, atol=1e-7, opset_versions=None, acceptable_error_percentage=None, ): opset_versions = opset_versions if opset_versions else [ 7, 8, 9, 10, 11, 12, 13, 14 ] for opset_version in opset_versions: self.opset_version = opset_version self.onnx_shape_inference = True onnx_test_common.run_model_test( self, model, input_args=inputs, rtol=rtol, atol=atol, acceptable_error_percentage=acceptable_error_percentage, ) if self.is_script_test_enabled and opset_version > 11: script_model = torch.jit.script(model) onnx_test_common.run_model_test( self, script_model, input_args=inputs, rtol=rtol, atol=atol, acceptable_error_percentage=acceptable_error_percentage, )
def test_nested_autograd(self): class Child(torch.autograd.Function): @staticmethod def forward(ctx, i): result = i.log() result_log = result.log() ctx.save_for_backward(result_log) return result_log @staticmethod def backward(ctx, grad_output): (result, ) = ctx.saved_tensors return grad_output * result class Parent(torch.autograd.Function): @staticmethod def forward(ctx, i): result_exp = i.exp() result_log = Child.apply(result_exp) ctx.save_for_backward(result_exp, result_log) return result_exp, result_log @staticmethod def backward(ctx, grad_output): (result, ) = ctx.saved_tensors return grad_output * result class Caller(torch.nn.Module): def forward(self, input): return Parent.apply(input) model = Caller() input = torch.ones(1, 5) run_model_test(self, model, input_args=(input, ))
def test_inline_with_scoped_tracing(self): class Exp(torch.autograd.Function): @staticmethod def forward(ctx, i): ctx.save_for_backward(input) return i.exp() @staticmethod def symbolic(g, input): return g.op("Exp", input) class LogLog(torch.autograd.Function): @staticmethod def forward(ctx, i): ctx.save_for_backward(input) return i.log().log() class Caller(torch.nn.Module): def forward(self, input): exp_result = Exp.apply(input) return LogLog.apply(exp_result) model = Caller() input = torch.ones(1) torch.jit._trace._trace_module_map = { _m: torch.typename(type(_m)) for _m in model.modules() } run_model_test(self, model, input_args=(input, )) torch.jit._trace._trace_module_map = None
def test_inline_and_symbolic(self): class Exp(torch.autograd.Function): @staticmethod def forward(ctx, i): ctx.save_for_backward(input) return i.exp() @staticmethod def symbolic(g, input): return g.op("Exp", input) class LogLog(torch.autograd.Function): @staticmethod def forward(ctx, i): ctx.save_for_backward(input) return i.log().log() class Caller(torch.nn.Module): def forward(self, input): exp_result = Exp.apply(input) return LogLog.apply(exp_result) model = Caller() input = torch.ones(1) run_model_test(self, model, input_args=(input, ))
def test_partial_output(self): class PartialOut(torch.autograd.Function): @staticmethod def forward(ctx, input): ctx.save_for_backward(input) values, indices = torch.topk(input, 3) return values class Caller(torch.nn.Module): def forward(self, input): return PartialOut.apply(input) model = Caller() input = torch.ones(1, 5) run_model_test(self, model, input_args=(input, ))
def test_register_custom_op(self): class MyClip(torch.autograd.Function): @staticmethod def forward(ctx, input, scalar): ctx.save_for_backward(input) return input.clamp(min=scalar) class MyRelu(torch.autograd.Function): @staticmethod def forward(ctx, input): ctx.save_for_backward(input) return input.clamp(min=0) class MyModule(torch.nn.Module): def __init__(self): super().__init__() self.clip = MyClip.apply self.relu = MyRelu.apply def forward(self, x): h = self.clip(x, 2) h = self.relu(h) return h def symbolic_pythonop(ctx: torch.onnx.SymbolicContext, g, *args, **kwargs): n = ctx.cur_node name = kwargs["name"] if name == "MyClip": return g.op("Clip", args[0], min_f=args[1], outputs=n.outputsSize()) elif name == "MyRelu": return g.op("Relu", args[0], outputs=n.outputsSize()) else: return symbolic_helper._unimplemented( "prim::PythonOp", "unknown node kind: " + name) from torch.onnx import register_custom_op_symbolic register_custom_op_symbolic("prim::PythonOp", symbolic_pythonop, 1) x = torch.randn(2, 3, 4, requires_grad=True) model = MyModule() onnx_test_common.run_model_test(self, model, input_args=(x, ))
def test_single_output(self): class SingleOut(torch.autograd.Function): @staticmethod def forward(ctx, i): result = i.exp() result = result.log() ctx.save_for_backward(result) return result @staticmethod def backward(ctx, grad_output): (result, ) = ctx.saved_tensors return grad_output * result class Caller(torch.nn.Module): def forward(self, input): result = input + 5 return SingleOut.apply(result) + 3 model = Caller() input = torch.ones(1) run_model_test(self, model, input_args=(input, ))
def test_symbolic(self): class MyClip(torch.autograd.Function): @staticmethod def forward(ctx, input, scalar): ctx.save_for_backward(input) return input.clamp(min=scalar) @staticmethod def symbolic(g, input, scalar): return g.op("Clip", input, min_f=scalar) class MyModule(torch.nn.Module): def __init__(self): super().__init__() self.clip = MyClip.apply def forward(self, x): h = self.clip(x, 2) return h x = torch.randn(2, 3, 4, requires_grad=True) model = MyModule() onnx_test_common.run_model_test(self, model, input_args=(x, ))