[ONNX] Export weight_norm #28618
[ONNX] Export weight_norm #28618
Conversation
torch/onnx/symbolic_opset9.py
Outdated
| # W = g * ((v) / ||v||) | ||
| # Compute norm_except_dim for l2 norm | ||
| axes = list(range(rank)) | ||
| axes.remove(dim) |
There was a problem hiding this comment.
We should support dim=None as well (computes a norm over the entire weight tensor).
Could you verify if axes.remove(dim) fails when dim=None, and add a test case with dim=None as well?
There was a problem hiding this comment.
Thanks Lara for mentioning that.
Looks like there is a conflicting behavior in torch.nn.utils.weight_norm regarding None dim.
Where dim could be negative for backwards axes, but when dim = None, it's overwitten to -1
pytorch/torch/nn/utils/weight_norm.py
Line 10 in 0c48092
So:
weight_norm(x, dim=-2) computes norm over dimension 'rank-2'
weight_norm(x, dim=-1) computes norm over all dimensions (equivalent of dim=None)
cc @houseroad to review the issue please.
This is not blocking the export currently. For now I've updated the symbolic to match torch function.
There was a problem hiding this comment.
@BowenBao Could you please take a look as well? Thanks.
facebook-github-bot
left a comment
facebook-github-bot
left a comment
There was a problem hiding this comment.
@houseroad has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
|
@pytorchbot rebase this please |
|
@houseroad I just rebased this PR. CI looks fine and it is approved. Could you please take a look? |
facebook-github-bot
left a comment
facebook-github-bot
left a comment
There was a problem hiding this comment.
@houseroad has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
|
@houseroad Can we merge this? Thanks! |
|
@houseroad merged this pull request in 7da11f4. |
6 participants
Export _weight_norm
Caffe2 tests are inplace
Looks like there is a conflicting behavior in torch.nn.utils.weight_norm regarding None dim.
Where dim could be negative for backwards axes, but when dim = None, it's overwitten to -1
pytorch/torch/nn/utils/weight_norm.py
Line 10 in 0c48092
For now, this symbolic to matches the current behavior. But this might need to be changed in the torch module.