Skip to content

ONNX Export LayerNorm #22265

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
lara-hdr wants to merge 6 commits into from
Closed

ONNX Export LayerNorm #22265

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
04efde0
ONNX Export LayerNorm
Jun 26, 2019
9c81b6f
remove dependency on dim
Jun 26, 2019
7b88ae6
merge and resolve conflicts
Jun 27, 2019
9d33fd6
merge and resolve conflicts + update op
Jul 1, 2019
7059042
use ATEN op for layer_norm with ONNX_ATEN_FALLBACK
Jul 1, 2019
21e6d43
add empty line
lara-hdr Jul 1, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
127 changes: 127 additions & 0 deletions test/onnx/expect/TestOperators.test_layer_norm_aten.expect
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,127 @@
ir_version: 4
producer_name: "pytorch"
producer_version: "1.1"
graph {
node {
input: "input"
input: "weight"
input: "bias"
output: "3"
op_type: "ATen"
attribute {
name: "cudnn_enable"
i: 1
type: INT
}
attribute {
name: "eps"
f: 1e-05
type: FLOAT
}
attribute {
name: "normalized_shape"
ints: 10
ints: 10
type: INTS
}
attribute {
name: "operator"
s: "layer_norm"
type: STRING
}
}
name: "torch-jit-export"
initializer {
dims: 10
dims: 10
data_type: 1
name: "bias"
raw_data: "\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000"
}
initializer {
dims: 10
dims: 10
data_type: 1
name: "weight"
raw_data: "\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?\000\000\200?"
}
input {
name: "input"
type {
tensor_type {
elem_type: 1
shape {
dim {
dim_value: 20
}
dim {
dim_value: 5
}
dim {
dim_value: 10
}
dim {
dim_value: 10
}
}
}
}
}
input {
name: "weight"
type {
tensor_type {
elem_type: 1
shape {
dim {
dim_value: 10
}
dim {
dim_value: 10
}
}
}
}
}
input {
name: "bias"
type {
tensor_type {
elem_type: 1
shape {
dim {
dim_value: 10
}
dim {
dim_value: 10
}
}
}
}
}
output {
name: "3"
type {
tensor_type {
elem_type: 1
shape {
dim {
dim_value: 20
}
dim {
dim_value: 5
}
dim {
dim_value: 10
}
dim {
dim_value: 10
}
}
}
}
}
}
opset_import {
version: 9
}
6 changes: 6 additions & 0 deletions test/onnx/test_operators.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -671,6 +671,12 @@ def forward(self, scores, bbox_deltas, im_info, anchors):
inputs = (scores, bbox_deltas, im_info, anchors)
self.assertONNX(model, inputs)

def test_layer_norm_aten(self):
model = torch.nn.LayerNorm([10, 10])
x = torch.randn(20, 5, 10, 10)
self.assertONNX(model, x,
operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK)
Copy link
Member

@houseroad houseroad Jul 1, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: one more empty line?



if __name__ == '__main__':
no_onnx_dep_flag = '--no-onnx'
Expand Down
24 changes: 24 additions & 0 deletions test/onnx/test_pytorch_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,30 @@ def wrapper(*args, **kwargs):
skipIfTravis = _skipper(lambda: os.getenv('TRAVIS'),
'Skip In Travis')

# skips tests for all versions below min_opset_version.
# if exporting the op is only supported after a specific version,
# add this wrapper to prevent running the test for opset_versions
# smaller than the currently tested opset_version
def skipIfUnsupportedMinOpsetVersion(min_opset_version):
def skip_dec(func):
def wrapper(self):
if self.opset_version < min_opset_version:
raise unittest.SkipTest("Skip verify test for unsupported opset_version")
return func(self)
return wrapper
return skip_dec

# skips tests for opset_versions listed in unsupported_opset_versions.
# if the caffe2 test cannot be run for a specific version, add this wrapper
# (for example, an op was modified but the change is not supported in caffe2)
def skipIfUnsupportedOpsetVersion(unsupported_opset_versions):
def skip_dec(func):
def wrapper(self):
if self.opset_version in unsupported_opset_versions:
raise unittest.SkipTest("Skip verify test for unsupported opset_version")
return func(self)
return wrapper
return skip_dec

def flatten(x):
return tuple(function._iter_filter(lambda o: isinstance(o, torch.Tensor))(x))
57 changes: 10 additions & 47 deletions test/onnx/test_pytorch_onnx_caffe2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@
import caffe2.python.onnx.backend as c2

from test_pytorch_common import skipIfTravis, skipIfNoLapack, skipIfNoCuda
from test_pytorch_common import skipIfUnsupportedOpsetVersion
import verify

skip = unittest.skip
Expand All @@ -52,31 +53,6 @@ def wrapper(self):
return func(self)
return wrapper

# skips tests for all versions below min_opset_version.
# if exporting the op is only supported after a specific version,
# add this wrapper to prevent running the test for opset_versions
# smaller than the currently tested opset_version
def skipIfUnsupportedMinOpsetVersion(min_opset_version):
def skip_dec(func):
def wrapper(self):
if self.opset_version < min_opset_version:
raise unittest.SkipTest("Skip verify test for unsupported opset_version")
return func(self)
return wrapper
return skip_dec

# skips tests for opset_versions listed in unsupported_opset_versions.
# if the caffe2 test cannot be run for a specific version, add this wrapper
# (for example, an op was modified but the change is not supported in caffe2)
def skipIfUnsupportedOpsetVersion(unsupported_opset_versions):
def skip_dec(func):
def wrapper(self):
if self.opset_version in unsupported_opset_versions:
raise unittest.SkipTest("Skip verify test for unsupported opset_version")
return func(self)
return wrapper
return skip_dec

# def import_model(proto, input, workspace=None, use_gpu=True):
# model_def = onnx.ModelProto.FromString(proto)
# onnx.checker.check_model(model_def)
Expand Down Expand Up @@ -140,7 +116,7 @@ def do_export(model, inputs, *args, **kwargs):
}


class TestCaffe2Backend(unittest.TestCase):
class TestCaffe2Backend_opset9(unittest.TestCase):
from torch.onnx.symbolic_helper import _export_onnx_opset_version
opset_version = _export_onnx_opset_version
embed_params = False
Expand Down Expand Up @@ -919,20 +895,6 @@ def forward(self, x):
x = torch.randn(*shape)
self.run_model_test(MyModel(), train=False, input=(x), batch_size=BATCH_SIZE, use_gpu=False)

def test_layer_norm(self):
shape = (20, 5, 10, 10)

class MyModel(torch.nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.ln = torch.nn.LayerNorm([5, 10, 10])

def forward(self, x):
return self.ln(x)

x = torch.randn(*shape)
self.run_model_test(MyModel(), train=False, input=(x,), batch_size=BATCH_SIZE, use_gpu=False)

def test_cosine_similarity(self):
shape = (100, 128)
x = torch.randn(*shape)
Expand Down Expand Up @@ -1993,7 +1955,7 @@ def f(self):
**extra_kwargs)

f.__name__ = test_name
setattr(TestCaffe2Backend, f.__name__, f)
setattr(TestCaffe2Backend_opset9, f.__name__, f)


def setup_rnn_tests():
Expand Down Expand Up @@ -2040,7 +2002,7 @@ def setup_rnn_tests():
test_count += 1

# sanity check that a representative example does exist
TestCaffe2Backend.test_gru_trilayer_forward_with_initial_state_without_sequence_lengths_with_dropout
TestCaffe2Backend_opset9.test_gru_trilayer_forward_with_initial_state_without_sequence_lengths_with_dropout

# make sure no one accidentally disables all the tests without
# noticing
Expand All @@ -2049,17 +2011,18 @@ def setup_rnn_tests():

# add the same test suite as above, but switch embed_params=False
# to embed_params=True
TestCaffe2BackendEmbed = type(str("TestCaffe2BackendEmbed"),
(unittest.TestCase,),
dict(TestCaffe2Backend.__dict__, embed_params=True))
TestCaffe2BackendEmbed_opset9 = type(str("TestCaffe2BackendEmbed_opset9"),
(unittest.TestCase,),
dict(TestCaffe2Backend_opset9.__dict__, embed_params=True))

# opset 10 tests
TestCaffe2Backend_opset10 = type(str("TestCaffe2Backend_opset10"),
(unittest.TestCase,),
dict(TestCaffe2Backend.__dict__, opset_version=10))
dict(TestCaffe2Backend_opset9.__dict__, opset_version=10))

TestCaffe2BackendEmbed_opset10 = type(str("TestCaffe2BackendEmbed_opset10"),
(unittest.TestCase,),
dict(TestCaffe2Backend.__dict__,
dict(TestCaffe2Backend_opset9.__dict__,
embed_params=True, opset_version=10))


Expand Down
16 changes: 12 additions & 4 deletions test/onnx/test_pytorch_onnx_onnxruntime.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,14 @@
import numpy as np
import io

from test_pytorch_onnx_caffe2 import skipIfUnsupportedMinOpsetVersion
from test_pytorch_common import skipIfUnsupportedMinOpsetVersion


class TestONNXRuntime(unittest.TestCase):
from torch.onnx.symbolic_helper import _export_onnx_opset_version
opset_version = _export_onnx_opset_version

def run_test(self, model, inputs):
def run_test(self, model, inputs, rtol=1e-05, atol=1e-08):
outputs = model(inputs)

# export the model to ONNX
Expand Down Expand Up @@ -50,10 +51,12 @@ def get_ort_inputs():
"number of outputs differ"

if isinstance(outputs, torch.Tensor):
assert np.allclose(get_numpy_value(outputs), ort_outs[0])
assert np.allclose(get_numpy_value(outputs), ort_outs[0],
rtol=rtol, atol=atol)
else :
for i in range(0, len(outputs)):
assert np.allclose(get_numpy_value_at_index(outputs, i), ort_outs[i])
assert np.allclose(get_numpy_value_at_index(outputs, i), ort_outs[i],
rtol=rtol, atol=atol)

def test_full_trace(self):
class FullModel(torch.nn.Module):
Expand Down Expand Up @@ -129,6 +132,11 @@ def forward(self, x):
x = torch.randn(1, 2, 3, 4, requires_grad=True)
self.run_test(MyModel(), x)

def test_layer_norm(self):
model = torch.nn.LayerNorm([10, 10])
x = torch.randn(20, 5, 10, 10)
self.run_test(model, x, rtol=1e-05, atol=1e-07)

def test_reduce_log_sum_exp(self):
class ReduceLogSumExpModel(torch.nn.Module):
def forward(self, input):
Expand Down
4 changes: 4 additions & 0 deletions torch/onnx/symbolic_helper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -240,6 +240,10 @@ def _set_opset_version(opset_version):
return
raise ValueError("Unsupported ONNX opset version: " + str(opset_version))

_operator_export_type = None
def _set_operator_export_type(operator_export_type):
global _operator_export_type
_operator_export_type = operator_export_type

# Metaprogram symbolics for each ATen native specialized cast operator.
# For e.g. we specify a function named `_cast_uint8_t` that instantiates an
Expand Down
33 changes: 27 additions & 6 deletions torch/onnx/symbolic_opset9.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -910,6 +910,33 @@ def batch_norm(g, input, weight, bias, running_mean, running_var, training, mome
return res


@parse_args('v', 'is', 'v', 'v', 'f', 'i')
def layer_norm(g, input, normalized_shape, weight, bias, eps, cudnn_enable):
Copy link
Member

@houseroad houseroad Jul 1, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you add a check here, if export mode is ONNX_ATEN_FALLBACK, we still export layer norm as ATen operator. Otherwise, the new export logic will significantly degrade the performance.

Reference:

pytorch/test/onnx/export_onnx_tests_generator.py

Line 104 in 173f224

operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK)

Copy link
Contributor Author

@lara-hdr lara-hdr Jul 1, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can add a function _set_operator_export_type (similar to _set_opset_version) in symbolic_helper.py, to save the operator_export_type and access it here

Copy link
Member

@houseroad houseroad Jul 1, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sounds good to me.

if sym_help._operator_export_type == torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK:
return g.op("ATen", input, weight, bias, normalized_shape_i=normalized_shape,
eps_f=eps, cudnn_enable_i=cudnn_enable, operator_s="layer_norm")

axes = [-i for i in range(len(normalized_shape), 0, -1)]

two_cst = g.op("Constant", value_t=torch.tensor(2.))
eps_cst = g.op("Constant", value_t=torch.tensor(eps))

mean = g.op("ReduceMean", input, axes_i=axes)
numerator = sub(g, input, mean)
# variance = e((x - e(x))^2), and (x - e(x)) is the numerator in the layer_norm formula
variance = g.op("ReduceMean", pow(g, numerator, two_cst), axes_i=axes)
denominator = sqrt(g, add(g, variance, eps_cst))

layer_norm = div(g, numerator, denominator)

if not (weight is None or weight.node().mustBeNone()):
layer_norm = mul(g, layer_norm, weight)
if not (bias is None or bias.node().mustBeNone()):
layer_norm = add(g, layer_norm, bias)

return layer_norm


@parse_args('v', 'v', 'v', 'v', 'v', 'i', 'f', 'f', 'i')
def instance_norm(g, input, weight, bias, running_mean, running_var, use_input_stats, momentum, eps, cudnn_enabled):
input_sizes = input.type().sizes()
Expand Down Expand Up @@ -973,12 +1000,6 @@ def type_as(g, self, other):
return g.op("ATen", self, other, operator_s="type_as")


@parse_args('v', 'is', 'v', 'v', 'f', 'i')
def layer_norm(g, self, normalized_shape, weight, bias, eps, cudnn_enable):
return g.op("ATen", self, weight, bias, normalized_shape_i=normalized_shape,
eps_f=eps, cudnn_enable_i=cudnn_enable, operator_s="layer_norm")


@parse_args('v', 'v', 'i', 'f')
def cosine_similarity(g, x1, x2, dim, eps):
return g.op("ATen", x1, x2, dim_i=dim, eps_f=eps, operator_s="cosine_similarity")
Expand Down
Loading