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

[复制链接] · 发表于 2020-4-27 16:58:20

本帖最后由 kkkaaa 于 2020-4-27 16:59 编辑

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



""" torch -> onnx -> rknn """



import torch

import numpy as np

from torch import nn





model_name = "little_model_func_conv"

ONNX_MODEL = model_name + '.onnx'

RKNN_MODEL = model_name + '_from_onnx' + '.rknn'



image_size_h = 2

image_size_w =1

num_channel =3 



# === torch 模型初始化 ===





class ToyModel(nn.Module):

    def __init__(self, ):

        super(ToyModel, self).__init__()

        self.op = nn.ReLU()



    def forward(self, x):

        x = self.op(x)

        return x



net = ToyModel()



print("==== network ====")

print(net)

net.eval()



# === 转化1: torch2onnx ===

print("--> torch model inference result")



input_tensor = torch.Tensor(np.arange(num_channel * image_size_h * image_size_w).reshape(1, num_channel, image_size_h, image_size_w))

torch_out = torch.onnx._export(net, input_tensor, ONNX_MODEL, export_params=True)



# === 转化2: onnx2rknn & rknn inference ===

from rknn.api import RKNN

rknn = RKNN()



print('--> Loading model')

ret = rknn.load_onnx(model=ONNX_MODEL)

if ret != 0:

    print('Load resnet50v2 failed!')

    exit(ret)

print('done')



# Build model

print('--> Building model')

ret = rknn.build(do_quantization=False, dataset='./dataset.txt')

if ret != 0:

    print('Build resnet50 failed!')

    exit(ret)

print('done')



# Export rknn model

print('--> Export RKNN model')

ret = rknn.export_rknn(RKNN_MODEL)

if ret != 0:

    print('Export resnet50v2.rknn failed!')

    exit(ret)

print('done')



# === rknn inference ===

# init runtime environment

print("--> Init runtime environment")

ret = rknn.init_runtime()

if ret != 0:

    print("Init runtime environment failed")

    exit(ret)

print('done')



# inference

print("--> Running rknn model")

rknn_input = input_tensor.numpy()[0].transpose(1,2,0)

print('----> rknn input')

print(rknn_input)

rknn_outputs = rknn.inference(inputs=[rknn_input])[0][0]  #[::-1]



# === torch inference ===

print('----> torch input')

print(input_tensor)

torch_inference_result = net(input_tensor)[0].detach().cpu().numpy()



# === onnx inference ===

import onnxruntime

def to_numpy(tensor):

    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()



ort_session = onnxruntime.InferenceSession(ONNX_MODEL)

ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(input_tensor)}

ort_outs = ort_session.run([x.name for x in ort_session.get_outputs()], ort_inputs)





# === compare & show results ===

print("~~~~~~ torch model infer output ~~~~~~")

print(torch_inference_result)

print("~~~~~~ onnx model infer output ~~~~~~")

print(ort_outs)

print("~~~~~~ rknn model infer output ~~~~~~")

print(rknn_outputs)

复制代码

输出结果如下

----> rknn input

[[[0. 2. 4.]]



 [[1. 3. 5.]]]

----> torch input

tensor([[[[0.],

          [1.]],



         [[2.],

          [3.]],



         [[4.],

          [5.]]]])

~~~~~~ torch model infer output ~~~~~~

[[[0.]

  [1.]]



 [[2.]

  [3.]]



 [[4.]

  [5.]]]

~~~~~~ onnx model infer output ~~~~~~

[array([[[[0.],

         [1.]],



        [[2.],

         [3.]],



        [[4.],

         [5.]]]], dtype=float32)]

~~~~~~ rknn model infer output ~~~~~~

[[[4.]

  [5.]]



 [[2.]

  [3.]]



 [[0.]

  [1.]]]
复制代码

只看该作者 · 发表于 2020-4-27 17:33:07

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

只看该作者 · 发表于 2020-4-27 17:45:23

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

只看该作者 · 发表于 2020-4-27 17:47:15

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

版本是 1.3.1b1

只看该作者 · 发表于 2020-4-27 17:49:15

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 的结果都一致了
但还是十分不解这是怎么回事

只看该作者 · 发表于 2020-4-27 20:37:24

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

只看该作者 · 发表于 2020-4-27 21:10:55

本帖最后由 kkkaaa 于 2020-4-27 21:13 编辑

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 的推断结果才一致

只看该作者 · 发表于 2020-5-14 21:51:54

本帖最后由 kkkaaa 于 2020-5-14 21:57 编辑

我知道原因了
因为我没写

rknn.config(reorder_channel="0 1 2")
复制代码

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