pytorch - ONNX 和 pytorch 的输出不同

Question

我尝试将我的 pytorch Resnet50 模型转换为 ONNX 并进行推理。转换程序没有错误，但是onnx模型的最终结果来自onnxruntime与pytorch的origin模型结果有较大差距。
可能的解决方案是什么？

ONNX 版本:1.5.0
pytorch版本:1.1.0
CUDA:9.0
系统:Ubuntu 18.06
Python:3.5

这里是转换代码

import torch
import models
from collections import OrderedDict
state_dict = "/home/yx-wan/newhome/workspace/filter-pruning-geometric-median/scripts/snapshots/resnet50-rate-0.7/best.resnet50.GM_0.7_76.82.pth.tar"
arch = 'resnet50'

def import_sparse(model,state_dict):
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k[7:]  # remove `module.`
        new_state_dict[name] = v
    model.load_state_dict(new_state_dict)
    print("sparse_model_loaded")
    return model

# initialize model
model = models.__dict__[arch](pretrained=False).cuda()

checkpoint = torch.load(state_dict)
model = import_sparse(model, checkpoint['state_dict'])
print("Top 1 precise of model: {}".format(checkpoint['best_prec1']))
dummy_input =torch.randn(1, 3, 224, 224).cuda()
torch.onnx.export(model, dummy_input, "{}.onnx".format(arch), verbose=True)

这是结果检查代码

import sys

from onnxruntime.datasets import get_example
import onnxruntime
import cv2
import numpy as np
import torch
import models
import onnxruntime
from collections import  OrderedDict
from my_tools import resize_img


def import_sparse(model,checkpoint):
    new_state_dict = OrderedDict()
    for k, v in checkpoint['state_dict'].items():
        name = k[7:]  # remove `module.`
        new_state_dict[name] = v
    model.load_state_dict(new_state_dict)
    return model

image_path = "./img652.jpg"
onnx_model_path = "/workplace/workspace/filter-pruning-geometric-median/resnet50.onnx"
ckpt="./scripts/snapshots/resnet50-rate-0.7/best.resnet50.GM_0.7_76.82.pth.tar"

img_ori = cv2.imread(image_path)  # BGR
img = cv2.cvtColor(img_ori, cv2.COLOR_BGR2RGB)
img, ratio_h, ratio_w = resize_img(img,224,224)
img = img - np.array([123.68, 116.78, 103.94],dtype=np.float32)
img_batch = np.expand_dims(img, 0)

# NHWC -> NCHW
img_batch = np.transpose(img_batch,[0,3,1,2])


example_model = get_example(onnx_model_path)
sess = onnxruntime.InferenceSession(example_model)
input_name = sess.get_inputs()[0].name
print("Input name  :", input_name)
input_shape = sess.get_inputs()[0].shape
print("Input shape :", input_shape)
input_type = sess.get_inputs()[0].type
print("Input type  :", input_type)
output_name = sess.get_outputs()[0].name
print("Output name  :", output_name)
output_shape = sess.get_outputs()[0].shape
print("Output shape :", output_shape)
output_type = sess.get_outputs()[0].type
print("Output type  :", output_type)

print("Input data shape{}".format(img_batch.shape))
assert(list(input_shape) == list(img_batch.shape))
result_onnx = sess.run([output_name], {input_name: img_batch})

# initialize model
model = models.__dict__["resnet50"]()
checkpoint = torch.load(ckpt,map_location='cpu')
best_prec1 = checkpoint['best_prec1']
model = import_sparse(model,checkpoint)
img_batch = torch.FloatTensor(img_batch)
with torch.no_grad():
    result_torch = model(img_batch)
    result_torch = result_torch.numpy()

print("max onnx-torch:{}".format(np.max(result_onnx-result_torch)))

检查代码的输出(带有一些警告，但我认为这并不重要)是

2019-08-09 02:59:21.378599853 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer4.2.bn3.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378654931 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer4.2.bn2.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378665235 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer4.2.bn1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378675069 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer4.1.bn3.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378686874 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer4.0.downsample.1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378698995 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer4.1.bn1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378718700 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.5.bn1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378729567 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.4.bn3.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378739657 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.4.bn2.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378752091 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.3.bn3.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378762533 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.3.bn1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378771168 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.2.bn3.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378781705 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.2.bn1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378792325 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.4.bn1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378802071 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.1.bn3.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378812061 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer4.0.bn3.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378822884 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.1.bn2.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378834198 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.0.downsample.1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378845176 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer1.2.bn1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378859324 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer2.0.bn1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378869709 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer2.0.bn2.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378883281 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.5.bn3.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378893302 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer2.3.bn1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378904876 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer1.1.bn2.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378915507 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer1.0.bn3.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378926638 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer4.0.bn2.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378938115 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.0.bn1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378948686 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'bn1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378958670 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.2.bn2.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378969125 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer4.1.bn2.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378979556 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer2.1.bn1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.378990553 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer2.2.bn1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379001126 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer1.2.bn2.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379011508 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer1.0.bn2.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379021900 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer1.0.downsample.1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379033504 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer2.2.bn3.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379044076 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer1.2.bn3.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379064049 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer1.1.bn1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379076654 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer2.0.bn3.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379089769 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer2.1.bn2.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379102140 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.0.bn2.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379114598 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.3.bn2.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379133520 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer2.2.bn2.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379144015 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer2.3.bn2.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379155771 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer1.1.bn3.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379167084 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer2.3.bn3.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379178303 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.0.bn3.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379189605 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer4.0.bn1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379199974 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.1.bn1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379211042 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer2.0.downsample.1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379221800 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer3.5.bn2.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379232566 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer1.0.bn1.num_batches_tracked'. It is not used by any node and should be removed from the model.
2019-08-09 02:59:21.379243442 [W:onnxruntime:Default, graph.cc:2263 CleanUnusedInitializers] Removing initializer 'layer2.1.bn3.num_batches_tracked'. It is not used by any node and should be removed from the model.
Input name  : 0
Input shape : [1, 3, 224, 224]
Input type  : tensor(float)
Output name  : 503
Output shape : [1, 1000]
Output type  : tensor(float)
Input data shape(1, 3, 224, 224)
max onnx-torch:104.89282989501953

Answer 1

通过在测试代码中运行 pytorch 模型推理之前添加 model.eval() 来解决问题。解决方案是 from the link

model = models.__dict__["resnet50"]()
checkpoint = torch.load(ckpt,map_location='cpu')
best_prec1 = checkpoint['best_prec1']
model = import_sparse(model,checkpoint)
model.eval()
img_batch = torch.FloatTensor(img_batch)
with torch.no_grad():
    result_torch = model(img_batch)
    result_torch = result_torch.numpy()