Skip to content
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

[PYTORCH]LayerNorm support added #5249

Merged
merged 1 commit into from
Apr 6, 2020
Merged

[PYTORCH]LayerNorm support added #5249

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
29 changes: 29 additions & 0 deletions python/tvm/relay/frontend/pytorch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -503,6 +503,34 @@ def _impl(inputs, input_types):
scale=scale)
return _impl

def _get_dims(data):
import torch
if isinstance(data, _expr.Expr):
dims = _infer_shape(data)
elif isinstance(data, list):
dims = data
elif isinstance(data, (torch.Tensor, np.ndarray)):
dims = data.shape
else:
msg = "Data type %s could not be parsed" % type(data)
raise AssertionError(msg)
return dims

def _layer_norm():
def _impl(inputs, input_types):
data = inputs[0]
ndims = len(_get_dims(inputs[1]))
assert ndims == 1, "Support only normalization over last one dimension."

return _op.nn.layer_norm(data,
gamma=inputs[1],
beta=inputs[2],
axis=-1,
Copy link
Member

Choose a reason for hiding this comment

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

it seems gamma is inputs[2] and beta is inputs[3]. @siju-samuel Can you confirm?
I get an error using this new conversion in my test. Above change fixed it.

Copy link
Member Author

@siju-samuel siju-samuel Apr 7, 2020 •
edited
Loading

Choose a reason for hiding this comment

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

@masahi Thanks a lot for pointing it out. Along with changing the inputs, need to update the center and scale flags to True. My testcase i was not initializing weights and thus it couldnt capure this mistake. #5257 will fix.

epsilon=float(inputs[4]),
center=False,
scale=False)
return _impl

def _transpose():
def _impl(inputs, input_types):
data = inputs[0]
Expand Down Expand Up @@ -1050,6 +1078,7 @@ def _wrap_const(c):
"aten::contiguous" : _contiguous(),
"aten::batch_norm" : _batch_norm(),
"aten::instance_norm" : _instance_norm(),
"aten::layer_norm" : _layer_norm(),
"aten::transpose" : _transpose(),
"aten::transpose_" : _transpose(),
"aten::t" : _transpose(),
Expand Down
4 changes: 4 additions & 0 deletions tests/python/frontend/pytorch/test_forward.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -561,6 +561,9 @@ def test_forward_instancenorm():
(torch.nn.InstanceNorm3d(16), inp_3d)]:
verify_model(ins_norm.eval(), input_data=inp)

def test_forward_layernorm():
inp = torch.rand((20, 5, 10, 10))
verify_model(torch.nn.LayerNorm(10).eval(), input_data=inp)

def test_forward_transpose():
torch.set_grad_enabled(False)
Expand Down Expand Up @@ -1132,6 +1135,7 @@ def forward(self, xs):
test_forward_contiguous()
test_forward_batchnorm()
test_forward_instancenorm()
test_forward_layernorm()
test_forward_transpose()
test_forward_size()
test_forward_view()
Expand Down