简单 op 通过pytorch->onnx->rknn 转化后推断结果很奇怪

以下是我实验的一个例子，转换 nn.Relu op，分别输出 torch, onnx, rknn 的推断结果。torch 和 onnx 一致，rknn 的结果和前两者是反的，也就是说 rknn_outputs = rknn.inference(inputs=[rknn_input])[0][0][::-1] 的话，三个模型的结果一致。
似乎和 channel 有关，但是还是有点奇怪

1. 
   
2. """ torch -> onnx -> rknn """
   
3. 
   
4. import torch
   
5. import numpy as np
   
6. from torch import nn
   
7. 
   
8. 
   
9. model_name = "little_model_func_conv"
   
10. ONNX_MODEL = model_name + '.onnx'
    
11. RKNN_MODEL = model_name + '_from_onnx' + '.rknn'
    
12. 
    
13. image_size_h = 2
    
14. image_size_w =1
    
15. num_channel =3
    
16. 
    
17. # === torch 模型初始化 ===
    
18. 
    
19. 
    
20. class ToyModel(nn.Module):
    
21.     def __init__(self, ):
    
22.         super(ToyModel, self).__init__()
    
23.         self.op = nn.ReLU()
    
24. 
    
25.     def forward(self, x):
    
26.         x = self.op(x)
    
27.         return x
    
28. 
    
29. net = ToyModel()
    
30. 
    
31. print("==== network ====")
    
32. print(net)
    
33. net.eval()
    
34. 
    
35. # === 转化1: torch2onnx ===
    
36. print("--> torch model inference result")
    
37. 
    
38. input_tensor = torch.Tensor(np.arange(num_channel * image_size_h * image_size_w).reshape(1, num_channel, image_size_h, image_size_w))
    
39. torch_out = torch.onnx._export(net, input_tensor, ONNX_MODEL, export_params=True)
    
40. 
    
41. # === 转化2: onnx2rknn & rknn inference ===
    
42. from rknn.api import RKNN
    
43. rknn = RKNN()
    
44. 
    
45. print('--> Loading model')
    
46. ret = rknn.load_onnx(model=ONNX_MODEL)
    
47. if ret != 0:
    
48.     print('Load resnet50v2 failed!')
    
49.     exit(ret)
    
50. print('done')
    
51. 
    
52. # Build model
    
53. print('--> Building model')
    
54. ret = rknn.build(do_quantization=False, dataset='./dataset.txt')
    
55. if ret != 0:
    
56.     print('Build resnet50 failed!')
    
57.     exit(ret)
    
58. print('done')
    
59. 
    
60. # Export rknn model
    
61. print('--> Export RKNN model')
    
62. ret = rknn.export_rknn(RKNN_MODEL)
    
63. if ret != 0:
    
64.     print('Export resnet50v2.rknn failed!')
    
65.     exit(ret)
    
66. print('done')
    
67. 
    
68. # === rknn inference ===
    
69. # init runtime environment
    
70. print("--> Init runtime environment")
    
71. ret = rknn.init_runtime()
    
72. if ret != 0:
    
73.     print("Init runtime environment failed")
    
74.     exit(ret)
    
75. print('done')
    
76. 
    
77. # inference
    
78. print("--> Running rknn model")
    
79. rknn_input = input_tensor.numpy()[0].transpose(1,2,0)
    
80. print('----> rknn input')
    
81. print(rknn_input)
    
82. rknn_outputs = rknn.inference(inputs=[rknn_input])[0][0]  #[::-1]
    
83. 
    
84. # === torch inference ===
    
85. print('----> torch input')
    
86. print(input_tensor)
    
87. torch_inference_result = net(input_tensor)[0].detach().cpu().numpy()
    
88. 
    
89. # === onnx inference ===
    
90. import onnxruntime
    
91. def to_numpy(tensor):
    
92.     return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
    
93. 
    
94. ort_session = onnxruntime.InferenceSession(ONNX_MODEL)
    
95. ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(input_tensor)}
    
96. ort_outs = ort_session.run([x.name for x in ort_session.get_outputs()], ort_inputs)
    
97. 
    
98. 
    
99. # === compare & show results ===
    
100. print("~~~~~~ torch model infer output ~~~~~~")
     
101. print(torch_inference_result)
     
102. print("~~~~~~ onnx model infer output ~~~~~~")
     
103. print(ort_outs)
     
104. print("~~~~~~ rknn model infer output ~~~~~~")
     
105. print(rknn_outputs)
     

复制代码
输出结果如下

1. ----> rknn input
   
2. [[[0. 2. 4.]]
   
3. 
   
4. [[1. 3. 5.]]]
   
5. ----> torch input
   
6. tensor([[[[0.],
   
7.           [1.]],
   
8. 
   
9.          [[2.],
   
10.           [3.]],
    
11. 
    
12.          [[4.],
    
13.           [5.]]]])
    
14. ~~~~~~ torch model infer output ~~~~~~
    
15. [[[0.]
    
16.   [1.]]
    
17. 
    
18. [[2.]
    
19.   [3.]]
    
20. 
    
21. [[4.]
    
22.   [5.]]]
    
23. ~~~~~~ onnx model infer output ~~~~~~
    
24. [array([[[[0.],
    
25.          [1.]],
    
26. 
    
27.         [[2.],
    
28.          [3.]],
    
29. 
    
30.         [[4.],
    
31.          [5.]]]], dtype=float32)]
    
32. ~~~~~~ rknn model infer output ~~~~~~
    
33. [[[4.]
    
34.   [5.]]
    
35. 
    
36. [[2.]
    
37.   [3.]]
    
38. 
    
39. [[0.]
    
40.   [1.]]]

复制代码

我还试过其他 op, 比如 nn.Conv2d, nn.ConvTranspose2d, 这些 op 就更看不出 rknn 结果和 torch 结果的联系

请确认rknn toolkit的版本，1.0.0以前版本输出是NHWC，1.0.0之后版本输出和原始模型一致。

leok 发表于 2020-4-27 17:45
请确认rknn toolkit的版本，1.0.0以前版本输出是NHWC，1.0.0之后版本输出和原始模型一致。 ...

版本是 1.3.1b1

leok 发表于 2020-4-27 17:45
请确认rknn toolkit的版本，1.0.0以前版本输出是NHWC，1.0.0之后版本输出和原始模型一致。 ...

我好像发现问题所在了
rknn_input = input_tensor.numpy()[0][::-1].transpose(1,2,0)
如果事先把 rknn_input 的 channel 维倒换过来，各种 op 的结果都一致了
但还是十分不解这是怎么回事

kkkaaa 发表于 2020-4-27 17:49
我好像发现问题所在了
rknn_input = input_tensor.numpy()[0][::-1].transpose(1,2,0)
如果事先把 rknn_in ...

rknn toolkit默认inference的data_format是nhwc

leok 发表于 2020-4-27 20:37
rknn toolkit默认inference的data_format是nhwc

嗯嗯，但是如果只做rknn_input = input_tensor.numpy()[0].transpose(1,2,0)，也就是只把 channel 维移到最后，结果还是颠倒的。

需要 rknn_input = input_tensor.numpy()[0][::-1].transpose(1,2,0)， 也就是先把 channel 维倒过来，再把 channel 维移到最后， onnx 和 rknn 的推断结果才一致

我知道原因了
因为我没写

1. rknn.config(reorder_channel="0 1 2")

复制代码
不指定 config 的话，当channel 个数等于3的时候看起来是默认 reorder_channel = "2 1 0"