def register_custom_nms_op():
# experimenting custom op registration.
from torch.onnx.symbolic_helper import parse_args
from torch.onnx.symbolic_opset9 import view, select, index_select, scatter
@parse_args('v', 'v', 'f', 'f', 'i')
def symbolic_nmsfilt(g, boxes, scores, iou_threshold, score_threshold, max_output_boxes):
# if should return all
if max_output_boxes <= 0:
max_output_boxes = 10000
shape = g.op("Shape", scores) # original shape
boxes = view(g, boxes, (1, -1, 4))
max_output_per_class = g.op('Constant', value_t=torch.tensor([max_output_boxes], dtype=torch.long))
iou_threshold = g.op('Constant', value_t=torch.tensor([iou_threshold], dtype=torch.float))
score_threshold = g.op('Constant', value_t=torch.tensor([score_threshold], dtype=torch.float))
# center_point_box == 1 is for our center_x, centr_y, width, height format
nms_out = g.op('NonMaxSuppression',
boxes, view(g, scores, (1, 1, -1)), max_output_per_class, iou_threshold, score_threshold,
center_point_box_i=1)
idx = view(g, select(g, nms_out, 1, g.op('Constant', value_t=torch.tensor([2], dtype=torch.long))), (-1,))
scores = view(g, scores, (-1,))
flat_shape = g.op("Shape", scores)
src = index_select(g, scores, 0, idx)
src = view(g, src, (-1,))
filt = g.op("ConstantOfShape", flat_shape)
filt = scatter(g, filt, 0, idx, src)
return view(g, filt, shape)
from torch.onnx import register_custom_op_symbolic
register_custom_op_symbolic('mtorch_ops::nmsfilt', symbolic_nmsfilt, 10)
def enable_custom_autograd_support():
# Initialize static objects needed to run custom autograd.Function's.
from onnxruntime.capi._pybind_state import (
register_forward_runner,
register_backward_runner,
unregister_python_functions,
)
from torch.onnx import register_custom_op_symbolic
from ._custom_autograd_function_exporter import _export
from ._custom_autograd_function_runner import call_python_forward_function, call_python_backward_function
from onnxruntime.training.ortmodule.torch_cpp_extensions import torch_interop_utils
import atexit
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
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 register_custom_nms_op():
# experimenting custom op registration.
from torch.onnx.symbolic_helper import parse_args
from torch.onnx.symbolic_opset9 import view, select
@parse_args('v', 'v', 'f', 'i')
def symbolic_nms(g, boxes, scores, iou_threshold, max_output_boxes):
# if should return all
if max_output_boxes <= 0:
max_output_boxes = 10000
boxes = view(g, boxes, (1, -1, 4))
max_output_per_class = g.op('Constant',
value_t=torch.tensor([max_output_boxes],
dtype=torch.long))
iou_threshold = g.op('Constant',
value_t=torch.tensor([iou_threshold],
dtype=torch.float))
# center_point_box == 1 is for our center_x, centr_y, width, height format
nms_out = g.op('NonMaxSuppression',
boxes,
view(g, scores, (1, 1, -1)),
max_output_per_class,
iou_threshold,
center_point_box_i=1)
idx = select(
g, nms_out, 1,
g.op('Constant', value_t=torch.tensor([2], dtype=torch.long)))
return view(g, idx, (-1, ))
from torch.onnx import register_custom_op_symbolic
register_custom_op_symbolic('mtorch_ops::nms', symbolic_nms, 10)
def register_custom_op():
def my_group_norm(g, input, num_groups, scale, bias, eps):
return g.op("mydomain::testgroupnorm", input, num_groups, scale, bias, epsilon_f=0.)
from torch.onnx import register_custom_op_symbolic
register_custom_op_symbolic("mynamespace::custom_group_norm", my_group_norm, 9)
def test_contrib_op_with_loop(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.gelu = torch.nn.GELU()
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):
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())
run_model_test(self, model, input=(x, ))
def test_custom_pythonop_pytorch(self):
# register_custom_op_symbolic(
# '<namespace>::inverse', my_inverse, <opset_version>)
register_custom_op_symbolic('::inverse', my_inverse, 1)
x = torch.randn(3, 3)
# Export model to ONNX
f = io.BytesIO()
torch.onnx.export(CustomInverse(), (x, ), f)
onnx_model = load(io.BytesIO(f.getvalue()))
self.assertIn('domain: "ai.onnx.contrib"', str(onnx_model))
model = CustomInverse()
pt_outputs = model(x)
so = _ort.SessionOptions()
so.register_custom_ops_library(_get_library_path())
# Run the exported model with ONNX Runtime
ort_sess = _ort.InferenceSession(f.getvalue(), so)
ort_inputs = dict((ort_sess.get_inputs()[i].name, input.cpu().numpy())
for i, input in enumerate((x, )))
ort_outputs = ort_sess.run(None, ort_inputs)
# Validate PyTorch and ONNX Runtime results
numpy.testing.assert_allclose(pt_outputs.cpu().numpy(),
ort_outputs[0],
rtol=1e-03,
atol=1e-05)
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 register_custom_op():
# experimenting custom op registration.
from torch.onnx.symbolic_helper import parse_args
from torch.onnx.symbolic_opset9 import select, unsqueeze, squeeze, _cast_Long
@parse_args('v', 'v', 'f')
def symbolic_multi_label_nms(
g, boxes, scores, iou_threshold
): #, max_output_per_class, iou_threshold, score_threshold):
boxes = unsqueeze(g, boxes, 0)
scores = unsqueeze(g, unsqueeze(g, scores, 0), 0)
max_output_per_class = g.op('Constant',
value_t=torch.tensor([2000],
dtype=torch.long))
iou_threshold = g.op('Constant',
value_t=torch.tensor([iou_threshold],
dtype=torch.float))
nms_out = g.op('NonMaxSuppression', boxes, scores,
max_output_per_class, iou_threshold)
return squeeze(
g,
select(
g, nms_out, 1,
g.op('Constant', value_t=torch.tensor([2], dtype=torch.long))),
1)
@parse_args('v', 'v', 'f', 'i', 'i', 'i')
def symbolic_roi_align(g, input, rois, spatial_scale, pooled_height,
pooled_width, sampling_ratio):
batch_indices = _cast_Long(
g,
squeeze(
g,
select(
g, rois, 1,
g.op('Constant',
value_t=torch.tensor([0], dtype=torch.long))), 1),
False)
rois = select(
g, rois, 1,
g.op('Constant',
value_t=torch.tensor([1, 2, 3, 4], dtype=torch.long)))
return g.op('RoiAlign',
input,
rois,
batch_indices,
spatial_scale_f=spatial_scale,
output_height_i=pooled_height,
output_width_i=pooled_width,
sampling_ratio_i=sampling_ratio)
from torch.onnx import register_custom_op_symbolic
register_custom_op_symbolic('DetectionHub::nms', symbolic_multi_label_nms,
10)
register_custom_op_symbolic('DetectionHub::roi_align_forward',
symbolic_roi_align, 10)
def enable_custom_autograd_support():
from onnxruntime.capi._pybind_state import register_forward_runner, register_backward_runner
from torch.onnx import register_custom_op_symbolic
from ._custom_autograd_function_exporter import _export
from ._custom_autograd_function_runner import call_python_forward_function, call_python_backward_function
register_forward_runner(call_python_forward_function)
register_backward_runner(call_python_backward_function)
register_custom_op_symbolic('::prim_PythonOp', _export, 1)
def register_custom_op():
# experimenting custom op registration.
from torch.onnx.symbolic import parse_args, _cast_Int, _cast_Long
def symbolic_multi_label_nms(g, boxes, scores,
max_output_per_class,
iou_threshold, score_threshold):
return g.op('NonMaxSuppression', boxes, scores,
max_output_per_class, iou_threshold,
score_threshold)
from torch.onnx import register_custom_op_symbolic
register_custom_op_symbolic('roi_ops::multi_label_nms',
symbolic_multi_label_nms)
def enable_custom_autograd_support():
from onnxruntime.capi._pybind_state import register_forward_runner, register_backward_runner, unregister_python_functions
from torch.onnx import register_custom_op_symbolic
from ._custom_autograd_function_exporter import _export
from ._custom_autograd_function_runner import call_python_forward_function, call_python_backward_function
import atexit
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)
register_custom_op_symbolic('::prim_PythonOp', _export, 1)
def register_custom_op():
"""
This function registers symbolic functions for
custom ops that are implemented as part of ONNX Runtime
"""
# Symbolic definition
def inverse(g, self):
return g.op("com.microsoft::Inverse", self).setType(self.type())
def gelu(g, self):
return g.op("com.microsoft::Gelu", self).setType(self.type())
def triu(g, self, diagonal):
return g.op("com.microsoft::Trilu", self, diagonal,
upper_i=1).setType(self.type())
def tril(g, self, diagonal):
return g.op("com.microsoft::Trilu", self, diagonal,
upper_i=0).setType(self.type())
# Op Registration
register_custom_op_symbolic('::inverse', inverse, _onnx_opset_version)
register_custom_op_symbolic('::gelu', gelu, _onnx_opset_version)
register_custom_op_symbolic('::triu', triu, _onnx_opset_version)
register_custom_op_symbolic('::tril', tril, _onnx_opset_version)
def test_register_aten_custom_op_symbolic(self):
self.addCleanup(unregister_custom_op_symbolic, "aten::gelu", 1)
def gelu(g, self):
return g.op("com.microsoft::Gelu", self).setType(self.type())
register_custom_op_symbolic("aten::gelu", gelu, 1)
model = torch.nn.GELU()
x = torch.randn(3, 3)
f = io.BytesIO()
torch.onnx.export(model, (x, ), f, opset_version=self.opset_version)
graph = onnx.load(io.BytesIO(f.getvalue()))
self.assertEqual(graph.graph.node[0].op_type, "Gelu")
self.assertEqual(graph.opset_import[1].domain, "com.microsoft")
def test_custom_opsets_gelu(self):
def gelu(g, self):
return g.op("com.microsoft::Gelu", self).setType(self.type())
register_custom_op_symbolic("::gelu", gelu, 1)
model = torch.nn.GELU()
x = torch.randn(3, 3)
f = io.BytesIO()
torch.onnx.export(model, (x, ), f,
opset_version=self.opset_version, custom_opsets={"com.microsoft": 1})
graph = onnx.load(io.BytesIO(f.getvalue()))
assert graph.graph.node[0].op_type == "Gelu"
assert graph.opset_import[0].version == self.opset_version
assert graph.opset_import[1].domain == 'com.microsoft'
assert graph.opset_import[1].version == 1
def register_custom_op():
def my_group_norm(g, input, num_groups, scale, bias, eps):
return g.op("mydomain::testgroupnorm",
input,
num_groups,
scale,
bias,
epsilon_f=0.)
from torch.onnx import register_custom_op_symbolic, set_custom_domain_version
# Optional: register custom domain version. If not registered, default version is 1
# set_custom_domain_version("mydomain", 2)
register_custom_op_symbolic("mynamespace::custom_group_norm",
my_group_norm, 9)
def register_custom_op():
"""
This function registers symbolic functions for
custom ops that are implemented as part of ONNX Runtime
"""
# Symbolic definition
def inverse(g, self):
return g.op("com.microsoft::Inverse", self)
def gelu(g, self):
return g.op("com.microsoft::Gelu", self)
# Op Registration
register_custom_op_symbolic('::inverse', inverse, _onnx_opset_version)
register_custom_op_symbolic('::gelu', gelu, _onnx_opset_version)
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 register_custom_op(is_ortmodule=False):
"""
This function registers symbolic functions for
custom ops that are implemented as part of ONNX Runtime
"""
# Symbolic definition
def inverse(g, self):
return g.op("com.microsoft::Inverse", self)
def gelu(g, self):
return g.op("com.microsoft::Gelu", self)
def triu(g, self, diagonal):
return g.op("com.microsoft::Trilu", self, diagonal, upper_i=1)
def tril(g, self, diagonal):
return g.op("com.microsoft::Trilu", self, diagonal, upper_i=0)
# Op Registration
register_custom_op_symbolic('::inverse', inverse, _onnx_opset_version)
register_custom_op_symbolic('::gelu', gelu, _onnx_opset_version)
register_custom_op_symbolic('::triu', triu, _onnx_opset_version)
register_custom_op_symbolic('::tril', tril, _onnx_opset_version)
if is_ortmodule:
@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()}'
'}')
return g.op("com.microsoft::ATenOp",
weight,
indices,
name_s='aten::embedding',
custom_attributes_json_s=custom_attributes_json)
register_custom_op_symbolic('::embedding', embedding,
_onnx_opset_version)
def test_custom_opsets_inverse(self):
class CustomInverse(torch.nn.Module):
def forward(self, x):
return torch.inverse(x) + x
def inverse(g, self):
return g.op("com.microsoft::Inverse", self).setType(self.type())
register_custom_op_symbolic("::inverse", inverse, 1)
model = CustomInverse()
x = torch.randn(2, 3, 3)
f = io.BytesIO()
torch.onnx.export(model, (x, ), f,
opset_version=self.opset_version, custom_opsets={"com.microsoft": 1})
graph = onnx.load(io.BytesIO(f.getvalue()))
self.assertEqual(graph.graph.node[0].op_type, "Inverse")
self.assertEqual(graph.opset_import[0].version, self.opset_version)
self.assertEqual(graph.opset_import[1].domain, "com.microsoft")
self.assertEqual(graph.opset_import[1].version, 1)
def test_custom_add(self):
op_source = """
#include <torch/script.h>
torch::Tensor custom_add(torch::Tensor self, torch::Tensor other) {
return self + other;
}
static auto registry =
torch::RegisterOperators("custom_namespace::custom_add", &custom_add);
"""
torch.utils.cpp_extension.load_inline(
name="custom_add",
cpp_sources=op_source,
is_python_module=False,
verbose=True,
)
class CustomAddModel(torch.nn.Module):
def forward(self, a, b):
return torch.ops.custom_namespace.custom_add(a, b)
def symbolic_custom_add(g, self, other):
return g.op("Add", self, other)
from torch.onnx import register_custom_op_symbolic
register_custom_op_symbolic(
"custom_namespace::custom_add", symbolic_custom_add, 9
)
x = torch.randn(2, 3, 4, requires_grad=False)
y = torch.randn(2, 3, 4, requires_grad=False)
model = CustomAddModel()
onnxir, _ = do_export(model, (x, y), opset_version=11)
onnx_model = onnx.ModelProto.FromString(onnxir)
prepared = c2.prepare(onnx_model)
caffe2_out = prepared.run(inputs=[x.cpu().numpy(), y.cpu().numpy()])
np.testing.assert_array_equal(caffe2_out[0], model(x, y).cpu().numpy())
def export_onnx_model(self, output_dir, verbose=False):
from torch.onnx import register_custom_op_symbolic
from torch.onnx.symbolic_helper import parse_args
import sys
torch_version = torch.__version__.split('.')
if int(torch_version[0]) == 1 and int(torch_version[1]) < 7:
NndctScreenLogger().error(
f'Only supprt exporting onnx model with pytorch 1.7 and later version'
)
return
@parse_args("v", "i", "i", "f", "i", "i", "i", "i")
def symbolic_fix_neuron(g, input, valmin, valmax, valamp, zero_point,
method, device_id, inplace):
#print(f'{valmax} {valamp} {method} {device_id}')
if valamp < sys.float_info.min:
scale = torch.tensor(sys.float_info.max).float(
) # Avoid exportor generating double type
else:
scale = torch.tensor(
1.0 /
valamp).float() # Avoid exportor generating double type
zero_point = torch.tensor(
0, dtype=torch.int8) # ONNX requires zero_point to be tensor
return g.op("DequantizeLinear",
g.op("QuantizeLinear", input, scale, zero_point),
scale, zero_point)
register_custom_op_symbolic("vai::fix_neuron", symbolic_fix_neuron, 9)
output_file = os.path.join(
output_dir, f"{self.quantizer.quant_model._get_name()}_int.onnx")
opset_version = torch.onnx.symbolic_helper._onnx_stable_opsets[-1]
device = GLOBAL_MAP.get_ele(NNDCT_KEYS.QUANT_DEVICE)
self.quantizer.reset_status_for_exporting()
model, input_args = to_device(self.quantizer.quant_model,
self._example_inputs, device)
torch.onnx.export(self.quantizer.quant_model,
input_args,
output_file,
verbose=verbose,
opset_version=opset_version)
def __init__(self,
filename: str,
image_size: int,
input_dtype: str = 'uint8',
n_channels=3,
n_batch=1,
opset_version=11,
**kwargs):
super(OnnxExporter, self).__init__(
image_size=image_size,
input_dtype=input_dtype,
filename=filename,
n_channels=n_channels,
n_batch=n_batch,
)
self.export_args = kwargs
self.opset_version = opset_version
self.export_args.update({'opset_version': opset_version})
"""
adapted from
https://github.com/pytorch/vision/blob/74679cc566f98398db13df0312cc11188733f1f3/torchvision/ops/_register_onnx_ops.py#L7
extend opset version
"""
opset_version = [9, 10]
for _onnx_opset_version in opset_version:
register_custom_op_symbolic('torchvision::nms',
type(self).symbolic_multi_label_nms,
_onnx_opset_version)
register_custom_op_symbolic('torchvision::roi_align',
type(self).roi_align,
_onnx_opset_version)
register_custom_op_symbolic('torchvision::roi_pool',
type(self).roi_pool,
_onnx_opset_version)
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(g, n, *args, **kwargs):
name = kwargs['name']
if name == "MyClip":
return g.op("Clip", args[0], min_f=args[1])
elif name == "MyRelu":
return g.op("Relu", args[0])
else:
return _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()
run_model_test(self, model, input=(x, ))
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)
def register_onnx_ops():
def convolution_sampler_ops(g, input1, input2, kH, kW, patchH, patchW,
padH, padW, dilationH, dilationW,
dilation_patchH, dilation_patchW, dH, dW):
return g.op("thales::correlation_layer",
input1,
input2,
kH_i=kH.node()['value'].item(),
kW_i=kW.node()['value'].item(),
patchH_i=patchH.node()['value'].item(),
patchW_i=patchW.node()['value'].item(),
padH_i=padH.node()['value'].item(),
padW_i=padW.node()['value'].item(),
dilationH_i=dilationH.node()['value'].item(),
dilationW_i=dilationW.node()['value'].item(),
dilation_patchH_i=dilation_patchH.node()['value'].item(),
dilation_patchW_i=dilation_patchW.node()['value'].item(),
dH_i=dH.node()['value'].item(),
dW_i=dW.node()['value'].item())
from torch.onnx import register_custom_op_symbolic
register_custom_op_symbolic("correlation_sampler::correlation_forward",
convolution_sampler_ops, 1)
def enable_custom_autograd_support():
# Initialize static objects needed to run custom autograd.Function's.
from onnxruntime.capi._pybind_state import register_forward_runner, register_backward_runner, unregister_python_functions
from torch.onnx import register_custom_op_symbolic
from ._custom_autograd_function_exporter import _export
from ._custom_autograd_function_runner import call_python_forward_function, call_python_backward_function
import atexit
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)
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
# For inference of the onnx model, you will need onnxruntime-gpu 1.6.0 (or nightly build).
import torch
import numpy as np
import argparse
import transformers
from torch.onnx import register_custom_op_symbolic
from torch.onnx.symbolic_helper import parse_args
from packaging import version
@parse_args('v', 'v', 'v', 'v','v', 'v', 'v', 'i', 'i')
def my_longformer_attention(g, input, weight, bias, mask, global_weight, global_bias, global_mask, num_heads, window):
return g.op("com.microsoft::LongformerAttention", input, weight, bias, mask, global_weight, global_bias, global_mask, num_heads_i=num_heads, window_i=window)
# namespace is onnxruntime which is registered in longformer_attention.cpp
register_custom_op_symbolic('onnxruntime::LongformerAttention', my_longformer_attention, 9)
# TODO: update the path according to output of "python setup.py install" when your python version is not 3.6
torch.ops.load_library(r'build/lib.linux-x86_64-3.6/longformer_attention.cpython-36m-x86_64-linux-gnu.so')
# mapping from model name to pretrained model name
MODELS = {
"longformer-base-4096": "allenai/longformer-base-4096",
"longformer-random-tiny": "patrickvonplaten/longformer-random-tiny" # A tiny model for debugging
}
is_debug = False
def parse_arguments():
parser = argparse.ArgumentParser()
def _register_custom_op():
from torch.onnx.symbolic_helper import parse_args, scalar_type_to_onnx
from torch.onnx.symbolic_opset9 import select, unsqueeze, squeeze, _cast_Long, reshape
@parse_args('v', 'v', 'f')
def symbolic_multi_label_nms(g, boxes, scores, iou_threshold):
boxes = unsqueeze(g, boxes, 0)
scores = unsqueeze(g, unsqueeze(g, scores, 0), 0)
max_output_per_class = g.op('Constant',
value_t=torch.tensor([sys.maxsize],
dtype=torch.long))
iou_threshold = g.op('Constant',
value_t=torch.tensor([iou_threshold],
dtype=torch.float))
nms_out = g.op('NonMaxSuppression', boxes, scores,
max_output_per_class, iou_threshold)
return squeeze(
g,
select(
g, nms_out, 1,
g.op('Constant', value_t=torch.tensor([2], dtype=torch.long))),
1)
@parse_args('v', 'v', 'f', 'i', 'i', 'i')
def roi_align(g, input, rois, spatial_scale, pooled_height, pooled_width,
sampling_ratio):
batch_indices = _cast_Long(
g,
squeeze(
g,
select(
g, rois, 1,
g.op('Constant',
value_t=torch.tensor([0], dtype=torch.long))), 1),
False)
rois = select(
g, rois, 1,
g.op('Constant',
value_t=torch.tensor([1, 2, 3, 4], dtype=torch.long)))
return g.op('RoiAlign',
input,
rois,
batch_indices,
spatial_scale_f=spatial_scale,
output_height_i=pooled_height,
output_width_i=pooled_width,
sampling_ratio_i=sampling_ratio)
@parse_args('v', 'v', 'f', 'i', 'i')
def roi_pool(g, input, rois, spatial_scale, pooled_height, pooled_width):
roi_pool = g.op('MaxRoiPool',
input,
rois,
pooled_shape_i=(pooled_height, pooled_width),
spatial_scale_f=spatial_scale)
return roi_pool, None
from torch.onnx import register_custom_op_symbolic
register_custom_op_symbolic('torchvision::nms', symbolic_multi_label_nms,
10)
register_custom_op_symbolic('torchvision::roi_align', roi_align, 10)
register_custom_op_symbolic('torchvision::roi_pool', roi_pool, 10)
