|
本帖最后由 kkkaaa 于 2020-4-24 11:12 编辑
模型的卷积用到 nn.ZeroPad2d, 似乎有问题,就单独拎出来测试转换。
torch -> onnx -> rknn (无量化): 可以转换成功,但是torch 和 rknn inference 结果不同,代码和结果都附在最后(虽然结果不一致,但是似乎有某种规律)
torch -> rknn(无量化):失败,
E Catch exception when loading pytorch model: little_model_zeropad.pt!
E Traceback (most recent call last):
E File "rknn/api/rknn_base.py", line 611, in rknn.api.rknn_base.RKNNBase.load_pytorch
E File "rknn/base/RKNNlib/app/importer/import_pytorch.py", line 97, in rknn.base.RKNNlib.app.importer.import_pytorch.ImportPytorch.run
E File "rknn/base/RKNNlib/converter/convert_pytorch.py", line 573, in rknn.base.RKNNlib.converter.convert_pytorch.convert_pytorch.__init__
E File "rknn/base/RKNNlib/converter/convert_pytorch.py", line 657, in rknn.base.RKNNlib.converter.convert_pytorch.convert_pytorch.model_simplify
E File "rknn/base/RKNNlib/converter/convert_pytorch.py", line 113, in rknn.base.RKNNlib.converter.convert_pytorch.torch_inference_engine.shape_pick
E File "rknn/base/RKNNlib/converter/convert_pytorch.py", line 148, in rknn.base.RKNNlib.converter.convert_pytorch.torch_inference_engine.__ir_shape_inference
E KeyError: 'aten::constant_pad_nd'
torch -> onnx -> rknn 的代码和输出结果
- class ZeroPadModel(nn.Module):
- def __init__(self):
- super(ZeroPadModel, self).__init__()
- self.static_padding = nn.ZeroPad2d((0,1,0,1))
- def forward(self, x):
- x = self.static_padding(x)
- return x
- # === torch 模型初始化 ===
- # net = Conv2dStaticSamePaddingModel(image_size)
- net = ZeroPadModel()
- print("==== network ====")
- print(net)
- net.eval()
- # === 转化1: torch2onnx ===
- print("--> torch model inference result")
- input_tensor = torch.rand(1,3, image_size, image_size)
- torch_out = torch.onnx._export(net, input_tensor, ONNX_MODEL, export_params=True)
- # === 转化2: onnx2rknn ===
- 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().transpose(0,2,3,1)
- rknn_outputs = rknn.inference(inputs=[rknn_input], data_format='nhwc')[0][0]
- # === torch inference ===
- torch_inference_result = net(input_tensor)[0].detach().cpu().numpy()
- # === compare & show results ===
- print("--> compare inference")
- print("input shape: ", input_tensor.shape)
- print("rknn input shape: ", rknn_input.shape)
- print("torch_inference shape: ", torch_inference_result.shape)
- print("rknn_outputs shape: ", rknn_outputs.shape)
- print("max error: ", np.max(torch_inference_result - rknn_outputs))
- print("~~~~~~ torch model infer output ~~~~~~")
- print(torch_inference_result)
- print("~~~~~~ rknn model infer output ~~~~~~")
- print(rknn_outputs)
输出:
~~~~~~ torch model infer output ~~~~~~
[[[0.05392462 0.12481469 0. ]
[0.44090217 0.44217068 0. ]
[0. 0. 0. ]]
[[0.86402655 0.47992283 0. ]
[0.72332394 0.12081265 0. ]
[0. 0. 0. ]]
[[0.92382944 0.56915957 0. ]
[0.28351116 0.7321489 0. ]
[0. 0. 0. ]]]
~~~~~~ rknn model infer output ~~~~~~
[[[0.9238281 0.56884766 0.56884766]
[0.28344727 0.7319336 0.7319336 ]
[0.28344727 0.7319336 0.7319336 ]]
[[0.86376953 0.47973633 0.47973633]
[0.72314453 0.12078857 0.12078857]
[0.72314453 0.12078857 0.12078857]]
[[0.05392456 0.12475586 0.12475586]
[0.44067383 0.44213867 0.44213867]
[0.44067383 0.44213867 0.44213867]]]
补充:onnx model 和 torch model 输出结果一致
|
|