pytorch->onnx-rknn 输入channel是否等于3决定推理结果是否正确

不知道是不是一个 bug
rknn 版本：1.3.1b1
输入给 rknn 的 tensor 已经做了 transpose, 满足 data_format = 'nhwc'
输出没有做任何处理，保持模型返回的结果


我测试的 toy model 定义如下

1. class ToyModel(nn.Module):
   
2.     def __init__(self, ):
   
3.         super(ToyModel, self).__init__()
   
4.         self.op = nn.ReLU()
   
5. 
   
6.     def forward(self, x):
   
7.         x = self.op(x)
   
8.         return x

复制代码
op 替换成其他简单 op 效果一样，为了明显看出规律，使用了 relu.


当输入 tensor 的形状为 (1,3, H, W) 时，结果的 channel 维颠倒了
当输入 tensor 的形状为 (1,C, H, W) and C!=3 时, 结果是正确的

以下是一个完整的可打印出 torcu, onnx, rknn 模型输入和推理结果的脚本

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)
     

复制代码这个问题困扰我好久了，请看看是不是有 bug, 谢谢

1. 试下1.3.2正式版本，
2. 文档有介绍，nchw都是要显式申明设置给rknn的，不写的话结果可能不符合预期。具体看下文档，我自己没玩过pytorch。

jefferyzhang 发表于 2020-5-1 22:24
1. 试下1.3.2正式版本，
2. 文档有介绍，nchw都是要显式申明设置给rknn的，不写的话结果可能不符合预期。具 ...

我试了 1.3.2, 而且显式指定 data_format = 'nhwc', 结果还是一样。
我觉得这本质上是一个从 onnx 到 rknn 转换的问题。

我把我观察到的现象描述的再具体一点。
当torch/onnx的输入是
tensor([[[[0.],
          [1.]],

         [[2.],
          [3.]]]])
即 rknn 的输入是
[[[0. 2.]]

[[1. 3.]]]时

rknn 输出是
[[[0.]
  [1.]]

[[2.]
  [3.]]]
----------
当 pytorch/onnx 输入是
tensor([[[[0.],
          [1.]],

         [[2.],
          [3.]],

         [[4.],
          [5.]]]])
即 rknn 输入是
[[[0. 2. 4.]]

[[1. 3. 5.]]] 时，
rknn 输出是
[[[4.]
  [5.]]

[[2.]
  [3.]]

[[0.]
  [1.]]]

这个问题挺困扰我的，虽然找到可以保持输出一致的方法，但是总觉得很奇怪

也不知道是不是我哪里做错了，不过我真的已经检查过很多次了

kkkaaa 发表于 2020-5-8 14:23
这个问题挺困扰我的，虽然找到可以保持输出一致的方法，但是总觉得很奇怪

也不知道是不是我哪里做错了，不 ...

显示指定NHWC或者NCHW，输入的数据也要是这个格式。如果不是的话，可以可以用np.transpose接口去转换。

找到原因了，和帖子 http://t.rock-chips.com/forum.php?mod=viewthread&tid=1542&extra= 是同一个问题