ONNX export only p-norms with p of 1 or 2

Raise code
@parse_args('v', 't', 'is', 'i')
def norm(g, self, p, dim, keepdim):
if p == 1:
f = _reduce_op_symbolic("ReduceL1")
elif p == 2:
f = _reduce_op_symbolic("ReduceL2")
else:
raise RuntimeError("ONNX export only p-norms with p of 1 or 2")
return f(g, self, dim=dim, keepdim=keepdim)
@parse_args('v', 'v', 'v', 'i')
def conv_tbc(g, input, weight, bias, pad):
if sym_help._operator_export_type == torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK:
return g.op("ATen", input, weight, bias, operator_s="conv_tbc", pad_i=pad)
Links to the raise (1)
https://github.com/pytorch/pytorch/blob/e56d3b023818f54553f2dc5d30b6b7aaf6b6a325/torch/onnx/symbolic_opset9.py#L1673Ways to fix
The function torch.norm is used to compute the matrix norm or vector norm of a given tensor.
Basic usage:
norm_order = 2
axis = 0
input = torch.tensor([[ 1, 2, 3],[-1, 1, 4]] , dtype= torch.float)
norm = torch.norm(input,
p = norm_order,
dim = axis,)
print(norm)
The parameter p
here is used to specify the order of norm being calculated.
i.e if p=1, L1-norm is calculated.
if p=2 it means to calculate the L2-norm.
The cause of this error is that Onnx conversion of a model that involves a norm operation with p
greater than 2 or p=0.
How to reproduce the error:
- Setup virtual environment
$ pip install --user pipenv
$ mkdir test_folder
$ cd test_folder
$ pipenv shell
- Install numpy and pytorch
pipenv install numpy torch
- Sample code
import torch
class NormModel(torch.nn.Module):
def forward(self, x):
# computing the norm of the tensor x to order p=3
x = torch.norm(x, p=3, dim=1)
x = x + 1
return x
model = NormModel()
dummy_input = torch.ones(2, 3, dtype=torch.float)
c = torch.tensor([[ 1, 2, 3],[-1, 1, 4]] , dtype= torch.float)
print(model(c))
with tempfile.TemporaryDirectory() as tmpdir:
model_path = os.path.join(tmpdir,"norm_model.onnx")
# export the sample model to onnx fomrmat
torch.onnx.export(model, (c), model_path,opset_version=10, verbose=True)
The ONNX export function only supports operation of matrix norm or vector norms that has order of 1 or 2. i.e L1-norm and L2 norms.
How to Fix This problem
In most machine learning projects the norms of order 1 and 2 are the most common cases. Therefore make sure why your model involves a norm different than p=1 and p=2.
If it is necessary that the norm outside the supported order is needed, then a custom computation of the norm should be implemented inside the model to bypass the pytorch operation( i.e torch.norm). This custom computation of matrix norm can be implemented using linalg.norm
Function.
import torch
import numpy as np
import os
import tempfile
class NormModel(torch.nn.Module):
def forward(self, x):
x = x.numpy # convert the x tensor to numpy
from numpy import linalg as LA # import the LA module to compute matrix norm
x = LA.norm(c, # input
ord=3, # this corrosponds to p in pytorch
axis=1)
# then finally convert the numpy to tensor bc the onnx export doesn't support numpy output
x =torch.from_numpy(x)
x = x + 1
return x
model = NormModel()
dummy_input = torch.ones(2, 3, dtype=torch.float)
c = torch.tensor([[ 1, 2, 3],[-1, 1, 4]] , dtype= torch.float)
print(model(c))
with tempfile.TemporaryDirectory() as tmpdir:
model_path = os.path.join(tmpdir,"norm_model.onnx")
# export the sample model to onnx fomrmat
torch.onnx.export(model, (c), model_path,opset_version=10, verbose=True)
print("\n\n successfuly converted to onnx and saved to norm_model.onnx\n\n")