ONNX转换中的Bug

在用rknn_toolkit=1.3.0转换ONNX时，发现对于ConvTranspose和PRelu存在尺寸判断错误。
测试使用代码如下：

1. import os
   
2. import onnx
   
3. from onnx import optimizer, shape_inference, version_converter, TensorProto
   
4. import onnxruntime
   
5. import numpy as np
   
6. import cv2
   
7. # import torch
   
8. from rknn.api import RKNN
   
9. np.random.seed(100)
   
10. # torch.manual_seed(100)
    
11. 
    
12. 
    
13. class ONNXModule(object):
    
14.     def __init__(self, name):
    
15.         super().__init__()
    
16.         self._name = name
    
17.         self._node_list, self._init_tensor, self._inpt_tensor, self._out_tensor = [], [], [], []
    
18.         self._graph = None
    
19.     @property
    
20.     def name(self):
    
21.         return self._name
    
22. 
    
23.     @property
    
24.     def model(self):
    
25.         inpt_value_info = []
    
26.         out_value_info = []
    
27.         init_value_data = []
    
28.         for init_t in self._inpt_tensor + self._init_tensor:
    
29.             inpt_value_info.append(onnx.helper.make_tensor_value_info(
    
30.                 init_t[0], init_t[2], tuple(init_t[1])))
    
31.         for init_t in self._init_tensor:
    
32.             init_value_data.append(onnx.helper.make_tensor(
    
33.                 init_t[0], init_t[2], tuple(init_t[1]),
    
34.                 np.random.randn(
    
35.                     *init_t[1]).astype(np.float32, order='C').flatten()))
    
36.         for init_t in self._out_tensor:
    
37.             out_value_info.append(onnx.helper.make_tensor_value_info(
    
38.                 init_t[0], init_t[2], tuple(init_t[1])))
    
39. 
    
40.         _graph = onnx.helper.make_graph(
    
41.             self._node_list, self._name,
    
42.             inpt_value_info, out_value_info, init_value_data
    
43.         )
    
44. 
    
45.         return onnx.helper.make_model(_graph)
    
46. 
    
47.     def __call__(self, inpt):
    
48.         pass
    
49. 
    
50. 
    
51. class M0(ONNXModule):
    
52.     # Good one
    
53.     def __init__(self):
    
54.         super().__init__('M0')
    
55.         # 1x3x10x10
    
56.         # self.conv = torch.nn.Conv2d(
    
57.         #     3, 13, kernel_size=(2, 2), stride=2)
    
58.         conv = onnx.helper.make_node(
    
59.             'Conv',
    
60.             inputs=['input', 'conv_W', 'conv_b'],
    
61.             outputs=['conv_1'],
    
62.             kernel_shape=[2, 2],
    
63.             strides=[2, 2]
    
64.         )
    
65.         # 1x13x5x5
    
66.         # self.relu = torch.nn.ReLU()
    
67.         relu = onnx.helper.make_node(
    
68.             'Relu',
    
69.             inputs=['conv_1'],
    
70.             outputs=['relu_1'],
    
71.         )
    
72.         self._node_list.extend([conv, relu])
    
73.         self._inpt_tensor.append(('input', [1, 3, 10, 10], TensorProto.FLOAT))
    
74.         self._init_tensor.append(('conv_W', [13, 3, 2, 2], TensorProto.FLOAT))
    
75.         self._init_tensor.append(('conv_b', [13], TensorProto.FLOAT))
    
76.         self._out_tensor.append(('relu_1', [1, 13, 5, 5], TensorProto.FLOAT))
    
77. 
    
78.     def __call__(self, inpt):
    
79.         pass
    
80.         # x = self.conv(inpt)
    
81.         # x = self.relu(x)
    
82.         # return x
    
83. 
    
84. 
    
85. class M1(ONNXModule):
    
86.     # Deconv Out Channel Error
    
87.     def __init__(self):
    
88.         super().__init__('M1')
    
89.         # 1x3x10x10
    
90.         # self.deconv = torch.nn.ConvTranspose2d(
    
91.         #     3, 13, kernel_size=(2, 2), stride=2)
    
92.         deconv = onnx.helper.make_node(
    
93.             'ConvTranspose',
    
94.             inputs=['input', 'deconv_W', 'deconv_b'],
    
95.             outputs=['deconv_1'],
    
96.             kernel_shape=[2, 2],
    
97.             strides=[2, 2]
    
98.         )
    
99.         # 1x13x20x20
    
100.         # self.relu = torch.nn.ReLU()
     
101.         relu = onnx.helper.make_node(
     
102.             'Relu',
     
103.             inputs=['deconv_1'],
     
104.             outputs=['relu_1'],
     
105.         )
     
106.         self._node_list.extend([deconv, relu])
     
107.         self._inpt_tensor.append(('input', [1, 3, 10, 10], TensorProto.FLOAT))
     
108.         self._init_tensor.append(
     
109.             ('deconv_W', [3, 13, 2, 2], TensorProto.FLOAT))
     
110.         self._init_tensor.append(('deconv_b', [13], TensorProto.FLOAT))
     
111.         self._out_tensor.append(('relu_1', [1, 13, 20, 20], TensorProto.FLOAT))
     
112. 
     
113.     def forward(self, inpt):
     
114.         pass
     
115.         # x = self.deconv(inpt)
     
116.         # x = self.relu(x)
     
117.         # return x
     
118. 
     
119. 
     
120. class M2(ONNXModule):
     
121.     # PReLU error
     
122.     def __init__(self):
     
123.         super().__init__('M2')
     
124.         # 1x3x10x10
     
125.         # self.conv = torch.nn.ConvTranspose2d(
     
126.         #     3, 6, kernel_size=(3, 3), stride=1, padding=1)
     
127.         conv = onnx.helper.make_node(
     
128.             'Conv',
     
129.             inputs=['input', 'conv_W', 'conv_b'],
     
130.             outputs=['conv_1'],
     
131.             kernel_shape=[3, 3],
     
132.             strides=[1, 1],
     
133.             pads=[1, 1, 1, 1]
     
134.         )
     
135.         # 1x6x10x10
     
136.         # self.prelu = torch.nn.PReLU(6)
     
137.         prelu = onnx.helper.make_node(
     
138.             'PRelu',
     
139.             inputs=['conv_1', 'slop_1'],
     
140.             outputs=['prelu_1']
     
141.         )
     
142.         self._node_list.extend([conv, prelu])
     
143. 
     
144.         self._inpt_tensor.append(('input', [1, 3, 10, 10], TensorProto.FLOAT))
     
145.         self._init_tensor.append(('conv_W', [6, 3, 3, 3], TensorProto.FLOAT))
     
146.         self._init_tensor.append(('conv_b', [6], TensorProto.FLOAT))
     
147.         self._init_tensor.append(('slop_1', [1, 6, 1, 1], TensorProto.FLOAT))
     
148.         self._out_tensor.append(('prelu_1', [1, 6, 10, 10], TensorProto.FLOAT))
     
149. 
     
150.     def forward(self, inpt):
     
151.         pass
     
152.         # x = self.conv(inpt)
     
153.         # x = self.prelu(x)
     
154.         # return x
     
155. 
     
156. class M3(ONNXModule):
     
157.     # Sub error
     
158.     def __init__(self):
     
159.         super().__init__('M3')
     
160.         # 1x3x10x10
     
161.         conv_1 = onnx.helper.make_node(
     
162.             'Conv',
     
163.             inputs=['input', 'conv_1W', 'conv_1b'],
     
164.             outputs=['conv_1'],
     
165.             kernel_shape=[2, 2],
     
166.             strides=[2, 2],
     
167.             pads=[0, 0, 0, 0]
     
168.         )
     
169.         # 1x3x10x10
     
170.         max_pool = onnx.helper.make_node(
     
171.             'MaxPool',
     
172.             inputs=['input'],
     
173.             outputs=['maxpool_1'],
     
174.             kernel_shape=[2, 2],
     
175.             strides=[2, 2]
     
176.         )
     
177.         sub = onnx.helper.make_node(
     
178.             'Sub',
     
179.             inputs=['maxpool_1', 'conv_1'],
     
180.             outputs=['result']
     
181.         )
     
182.         self._node_list.extend([conv_1, max_pool, sub])
     
183. 
     
184.         self._inpt_tensor.append(('input', [1, 3, 10, 10], TensorProto.FLOAT))
     
185.         self._init_tensor.append(('conv_1W', [3, 3, 2, 2], TensorProto.FLOAT))
     
186.         self._init_tensor.append(('conv_1b', [3], TensorProto.FLOAT))
     
187.         self._out_tensor.append(('result', [1, 3, 5, 5], TensorProto.FLOAT))
     
188. 
     
189.     def forward(self, inpt):
     
190.         pass
     
191. 
     
192. def export_onnx(inst_, input_, name):
     
193.     output = inst_(input_)
     
194.     onnx_m = inst_.model
     
195.     onnx_m = shape_inference.infer_shapes(onnx_m)
     
196.     onnx_m = optimizer.optimize(onnx_m)
     
197.     onnx.checker.check_model(onnx_m)
     
198.     onnx.save(onnx_m, name)
     
199. 
     
200. 
     
201. def export_rknn(onnx_, rknn_, dataset_f='./dataset.txt'):
     
202.     assert os.path.exists(onnx_), f"onnx file '{onnx_}' not exist"
     
203.     rknn = RKNN()
     
204.     rknn.config(
     
205.         channel_mean_value='127.5 127.5 127.5 127.5', reorder_channel='0 1 2',
     
206.         quantized_dtype='dynamic_fixed_point-16',
     
207.         epochs=1000)
     
208.     rknn.load_onnx(model=onnx_)
     
209.     rknn.build(do_quantization=True, dataset=dataset_f)
     
210.     rknn.export_rknn(rknn_)
     
211.     rknn.release()
     
212. 
     
213. 
     
214. def run_rknn(rknn_, onnx_, inpt):
     
215.     rknn = RKNN()
     
216.     rknn.config(
     
217.         channel_mean_value='127.5 127.5 127.5 127.5', reorder_channel='0 1 2')
     
218.     rknn.load_rknn(rknn_)
     
219. 
     
220.     rknn.init_runtime()
     
221.     rk_out = rknn.inference([inpt], data_format='nchw')
     
222. 
     
223.     t_inpt = inpt.astype(np.float32, order='C')
     
224.     t_inpt -= np.array([127.5, 127.5, 127.5], dtype=np.float32).reshape(1, 3, 1, 1)
     
225.     t_inpt /= 127.5
     
226.     # with torch.no_grad():
     
227.     #     th_out = inst_(t_inpt)
     
228.     onrt = onnxruntime.InferenceSession(onnx_, None)
     
229.     th_out = onrt.run([], {'input': t_inpt})
     
230.     if isinstance(th_out, (list, tuple)):
     
231.         pass
     
232.     else:
     
233.         th_out = [th_out]
     
234.     # th_out = [t.cpu().numpy() for t in th_out]
     
235.     for rk_o, th_o in zip(rk_out, th_out):
     
236.         print(rk_o.shape, th_o.shape)
     
237.         assert rk_o.size == th_o.size, "size mismatch ({rk_o.size}) ({th_o.size})"
     
238.         rk_o = rk_o.reshape(th_o.shape)
     
239.         print('abs mean: ', np.abs(rk_o - th_o).mean())
     
240. 
     
241.     # rknn.release()
     
242. 
     
243. 
     
244. DATASET_NUM = 1
     
245. # dataset_file = open('Data/TOM/')
     
246. all_img = np.random.randint(0, high=256, size=(
     
247.     DATASET_NUM, 10, 10, 3), dtype=np.uint8)
     
248. dataset_f = open('./dataset.txt', 'w')
     
249. os.makedirs('input', exist_ok=True)
     
250. for i in range(DATASET_NUM):
     
251.     img = all_img[i]
     
252.     cv2.imwrite(f'input/{i}.png', img)
     
253.     dataset_f.write(f'input/{i}.png\n')
     
254. dataset_f.close()
     
255. 
     
256. test_data = cv2.imread('input/0.png')
     
257. test_data = cv2.cvtColor(test_data, cv2.COLOR_BGR2RGB)
     
258. test_data = np.transpose(
     
259.     test_data[None, :, :, :], (0, 3, 1, 2)).copy()  # input used for test
     
260. # input_data = torch.from_numpy(test_data).float()  # input used for export
     
261. input_data = test_data
     
262. 
     
263. 
     
264. M0_inst = M0()
     
265. M1_inst = M1()
     
266. M2_inst = M2()
     
267. # M3_inst = M3()
     
268. m_l = [M0_inst, M1_inst, M2_inst]
     
269. n_l = [m.name for m in m_l]
     
270. ox_l = [m.name+'.onnx' for m in m_l]
     
271. rk_l = [m.name+'.rknn' for m in m_l]
     
272. 
     
273. for m, ox in zip(m_l, ox_l):
     
274.     export_onnx(m, input_data, ox)
     
275. 
     
276. for ox, rk, m, n in zip(ox_l, rk_l, m_l, n_l):
     
277.     print(f"=====start {n}========")
     
278.     try:
     
279.         export_rknn(ox, rk)
     
280.         print(f"{n} export_rknn ok")
     
281.         run_rknn(rk, ox, test_data)
     
282.         print(f"{n} run_rknn ok")
     
283.     except (Exception, TypeError, NameError) as ex:
     
284.         print(ex)
     
285.         print(f"{n} failed")
     
286.     finally:
     
287.         print(f"======end {n}=========")
     
288.         print()
     

复制代码在M1中使用ConvTranpose，在M2中使用PRelu，分别报如下错误
M1:
E ValueError: output_shape does not match filter's output channels, 3 != 13

M2:
E ValueError: Dimensions must be equal, but are 6 and 10 for 'PRelu_prelu_1_1_1/mul' (op: 'Mul') with input shapes: [1,6,1,1], [1,10,10,6].

这个onnx，有尝试过用onnx自己加载并推理成功么？

jefferyzhang 发表于 2020-5-20 14:30
这个onnx，有尝试过用onnx自己加载并推理成功么？

这个里面使用了onnxruntime计算真值，如果build时报错或量化后误差较大应该都是有问题的
每个模型onnxruntime运行都是没有问题的，可以单独把run_rknn里面的onnxruntime代码取出来进行测试

使用的环境是python 3.6
onnxruntime==1.2.0
onnx==1.4.1
rknn_toolkit==1.3.2

问题已经上报了

NPU Team回复：253412
M2的问题后续 RKNN Toolkit 1.3.3 版本会修复。
M1的问题后续需要更新 RKNN Toolkit 1.3.3 和 最新的npu驱动才可修复。