如何将割炬层的标识符传递给ONNX图

问题描述

对于TensorFlow，层的名称也是ONNX图中节点的name：

import onnx
from tensorflow import keras
import keras2onnx

model = keras.Sequential([
    keras.layers.Conv2D(4,(3,3),input_shape=(5,5,2),name='rainbow'),keras.layers.Conv2D(8,name='flower')])

model_onnx = keras2onnx.convert_keras(model)
nodes = model_onnx.graph.node

for node in nodes:
    print('type:',node.op_type,'- name:',node.name)

输出：

type: Conv - name: rainbow
type: Conv - name: flower

但是，对于PyTorch，名称字段是一个空字符串：

import torch
from torch import nn
import onnx

class NetConv(nn.Module):
    def __init__(self):
        super(NetConv,self).__init__()
        self.conv_banana = nn.Conv2d(2,4,3)
        self.conv_apple = nn.Conv2d(4,8,3)

    def forward(self,x):
        x = self.conv_banana(x)
        return self.conv_apple(x)

model = NetConv()
image = torch.randn(1,2,5)

filename_model = 'model.onnx'
torch.onnx.export(model,image,filename_model,export_params=False)
model_onnx = onnx.load(filename_model)
nodes = model_onnx.graph.node

for node in nodes:
    print('type:',node.name)
    print('inputs:',node.input)

输出：

type: Conv - name: 
inputs: ['input.1','conv_banana.weight','conv_banana.bias']
type: Conv - name: 
inputs: ['5','conv_apple.weight','conv_apple.bias']

我希望将图层名称分配给图节点的name属性，但是此属性为''。

是否有任何干净的方法来获取标识符，以便可以将PyTorch中的图层分配给ONNX图中的节点？

我唯一想到的解决方法是从图中节点的input中提取层的名称，但是这种做法很麻烦。

更新：

上面的输出是使用Torch == 1.4.0生成的。 对于torch == 1.6.0，name属性不再是一个空字符串，而且也比一个空字符串更有用：

type: Conv - name: Conv_0
inputs: ['input.1','conv_banana.bias']
type: Conv - name: Conv_1
inputs: ['5','conv_apple.bias']

解决方法

暂无找到可以解决该程序问题的有效方法，小编努力寻找整理中！

如果你已经找到好的解决方法，欢迎将解决方案带上本链接一起发送给小编。

小编邮箱:dio#foxmail.com (将#修改为@）