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, ))
