I used the method of “torch.onnx.export” to convert into onnx format like the below codes
intermediate tensor name was given by the sequence number.
Is there a way to give a name string instead of a auto generated sequence?
class ExModel(nn.Module):
def __init__(self):
super(ExModel, self).__init__()
out_planes = 6
self.conv = nn.Sequential(
nn.Conv2d(3, out_planes, 3, 1, 1, groups=1, bias=False),
nn.BatchNorm2d(out_planes),
nn.ReLU(inplace=True)
)
def forward(self, x):
return self.conv(x)
def convert_onnx_named(m):
input_names = ["input_img", "in_1", "in_2"]
output_names = ["o1"]
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(m, dummy_input, "exmodel.onnx", verbose=True,
input_names=input_names, output_names=output_names)
m = ExModel()
convert_onnx_named(m)
the output of converting into onnx.
graph(%input_img : Float(1, 3, 224, 224),
%in_1 : Float(6, 3, 3, 3),
%in_2 : Float(6),
%conv.1.bias : Float(6),
%conv.1.running_mean : Float(6),
%conv.1.running_var : Float(6),
%conv.1.num_batches_tracked : Long()):
%7 : Float(1, 6, 224, 224) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1]](%input_img, %in_1), scope: ExModel/Sequential[conv]/Conv2d[0] # /root/.pyenv/versions/3.7.1/lib/python3.7/site-packages/torch/nn/modules/conv.py:340:0
%8 : Float(1, 6, 224, 224) = onnx::BatchNormalization[epsilon=1e-05, momentum=0.9](%7, %in_2, %conv.1.bias, %conv.1.running_mean, %conv.1.running_var), scope: ExModel/Sequential[conv]/BatchNorm2d[1] # /root/.pyenv/versions/3.7.1/lib/python3.7/site-packages/torch/nn/functional.py:1656:0
%o1 : Float(1, 6, 224, 224) = onnx::Relu(%8), scope: ExModel/Sequential[conv]/ReLU[2] # /root/.pyenv/versions/3.7.1/lib/python3.7/site-packages/torch/nn/functional.py:911:0
return (%o1)
I want to give a name ‘in_1’ for ‘%7’ and ‘in_2’ for ‘%8’.