解决UNEt pytorch模型转ONNX模型之后推理结果错误的问题

1 前言：

1.1  我用UNet模型来检测表格的行列线，模型（基于pytorch）训练好之后预测准确，想转换成ONNX模型来部署，结果遇到了转换后的ONNX模型推理结果有误的问题

2 问题排查：

  2.1  输入图片、预处理以及后处理是否一致

      pytorch模型下推理输入的是图片的tensor，输出保存的也是4维tensor；ONNX模型推理输入的也是图片的像素array，输出array通过cv2.imwrite进行保存

  图片预处理是在模型输入前进行，模型后处理sigmoid被封装进ONNX模型

   排查结果：

        ①pytorch模型换成array进行cv2保存，结果也是准确的 

　　 ② 输入的图片像素大致也是一样的

      ③ 添加图片后处理模块 WrappedModel，封装进ONNX模型

class WrappedModel(torch.nn.Module):
    def __init__(self,model):
        super().__init__()
        self.model =model

    def forward(self,x):#前向传播函数重写父类torch.nn.Module的子函数，只有添加forward函数，才能进行   y=WrappedModel的实例（x）
        print("增加sigmoid后处理")
        outs=self.model(x)
        new_outs=torch.sigmoid(outs)
        return new_outs

    def resume(self, model, path):
        if not os.path.exists(path):
            print("Checkpoint not found:" + path)
            return
        states = torch.load(path, map_location=self.device)#直接’cpu"??
        model.load_state_dict(states["state_dict"],strict=False)#states有两个key_value"state_dict","optimizer"
        model_sig = WrappedModel(model)
        print("Resume from " + path)
        return model_sig

 

  2.2  模型加载是否一致

       ①模型转换时确保(当初因为加载的states的全部参数导致结果错误):

　states = torch.load(path, map_location=self.device)#map_loaction: cpu 这里设置成cpu是为了让模型通过CPU计算

　##states是dict，有2个K-V对，一个是state_dict ,一个是optimizer，这里只需要加载state——dict(模型权重系数)
　model.load_state_dict(states["state_dict"],strict=False)

　　 ②初始化tensor类型：cpu or cuda以及设置默认tensor类型　

self.device = 'cpu'#torch.device('cpu')都行
torch.set_default_tensor_type('torch.FloatTensor')   

　　 ③：构建的网络也要设置计算设备（与之前设置的device统一）

model = UNET(in_channels=in_channels, out_channels=out_channels).to(self.device)

　　④：调用 onnx.export之前一定要设置model.eval()模式

(非训练模式下，BN层，dropout层等用于优化训练的层会被关闭），不会影响推理结果；

　　　　通常配合with torch.no_grad()一起使用

  2.3. 输入图片计算精度是否一致，输出结果保存数据一致

      图片标签json文件里也记录了图片像素值，pytorch预测的时候输入的是json文件读取到的像素值，确保精度一致；

     cv2.imread读取的是BGR格式、精度为int的图片，最好统一转换成RGB格式且精度为float32，再进行计算处理

3、最终解决：

　　见2.2

4、解析torch.device  torch.load()  model.load_state_dict()

参考https://blog.csdn.net/qq_28949847/article/details/129400579?ops_request_misc=&request_id=&biz_id=102&utm_term=torch.load_state_dict&utm_medium=distribute.pc_search_result.none-task-blog-2~all~sobaiduweb~default-0-129400579.nonecase&spm=1018.2226.3001.4187

https://blog.csdn.net/baidu_41879652/article/details/118307330

4.1 torch.device()

torch.device 代表torch.Tensor分配到的设备对象

torch.device 包含一个设备类型（'cpu'或者'cuda'设备类型）和可选的设备序号，如果设备序号不存在，则为当前设备

Tensor的设备可以通过tensor.device访问属性

4.1.1通过字符串构造设备

>>> device= torch.device('cpu')
>>> device
device(type='cpu')
>>> device=torch.device('cuda')
>>> device
device(type='cuda')

4.1.2 通过字符串+序号 构造设备

>>> device=torch.device('cuda',0)
>>> device
device(type='cuda', index=0)
>>> device=torch.device('cpu',0)
>>> device
device(type='cpu', index=0)
>>> 

4.1.3  直接用字符串构建设备

#字符串好像是在预先的集合里面特定的，不是随意更改的，如下所示，设备不止一个的时候加上 “：序号”

Expected one of：

{cpu, cuda, ipu, xpu, mkldnn, opengl, opencl, ideep, hip, ve, ort, mps, xla, lazy, vulkan, meta, hpu}

cuda1 = torch.device('cuda:0')
>>> torch.randn((2,3),device=cuda1)
tensor([[-0.0963,  0.8858, -0.4105],
        [-0.4756, -0.6081,  0.0584]], device='cuda:0')
>>> torch.randn((2,3),device='cpu')#先构建device再赋值，也可以直接对tensor的device赋值特定字符串： ‘cpu’ 'cuda:0'（这些特定字符串可能是默认的device）
tensor([[ 1.2358, -0.9411, -1.2635],
        [-1.6386,  0.2647, -0.4192]])
>>> cpu = torch.device('cpu')
>>> torch.randn((2,3),device=cpu)
tensor([[ 1.7858,  1.5141, -0.2823],
        [ 0.6104, -0.0601, -0.5572]])
>>> torch.randn((2,3),device='cuda:0')
tensor([[-0.1011, -0.9072,  0.2595],
        [ 1.0570, -0.7164,  0.1800]], device='cuda:0')

4.1.4 通过单个序列号构建设备

>>> torch.randn((2,3),device = torch.device('cuda:0'))
tensor([[ 0.2886, -0.4556,  0.8388],
        [-0.2063,  1.1761,  1.9339]], device='cuda:0')
>>> torch.randn((2,3),device = torch.device('cpu'))
tensor([[-0.1918,  0.3366,  0.6095],
        [-1.7267,  1.1189, -1.9839]])
>>> torch.randn((2,3),device = 'cuda:0')
tensor([[-0.2101,  0.9385,  0.7495],
        [ 0.2342,  2.0308,  0.6125]], device='cuda:0')
>>> torch.randn((2,3),device = 0)
tensor([[0.7180, 0.1004, 0.4057],
        [0.4966, 0.1925, 0.9119]], device='cuda:0')

4.2 torch.load()

torch.load(model_path, map_location=None,pickle_module=pickle, **pickle_load_args)

一般只使用前面2个参数：map_location参数用于重定向，此前模型参数在CPU中，现在要加载到cuda:0上；或者将cuda：0中的模型加载到cuda1

4.2.3 map_location = None

不指定map_location，默认以训练保存模型时的位置加载

model = Model(base_channel=32, num_joints=17)
weights_dict = torch.load("xxxx.pth")
# 打印模型权重所在的位置
print(weights_dict['conv1.weight'].device)
print('weights_dict.keys():', weights_dict.keys())

4.2.4 map_location = 'cpu'

4.2.5 map_location={xx:xx}

model =Model(input_channel, output_channels）2 print('model:', model)
weights_dict = torch.load("./xxxx.pth", map_location={'cuda:0':'cpu'})
print('weights_dict:', weights_dict)
# 打印模型权重所在的位置
print(weights_dict['conv1.weight'].device)
print('weights_dict.keys():', weights_dict.keys())

模型加载从cuda：0变成cpu，要是目标device不存在，则map_location设置无效，还是保持原先device

4.3 torch.load_state_dict()
在pytorch构建好模型网络结构之后，需要加载模型权重

# 模型初始化
model = Model(input_channel, output_channels)
# 读取官方的模型参数
weights_dict = torch.load("./xxxxxx.pth", map_location='cpu')
# 加载官方模型参数到模型中
model.load_state_dict(weights_dict, strict=False)

ps：weights_dict中有state_dict和optimizer，需要鉴别，测试的时候只需要加载state_dict

strict=True: 要求加载的训练权重层数的键值与新构建的模型的权重层数完全吻合，否则报错：key对应不上

strict=False:查询训练权重中与新构建的网络中匹配层的键值就行，没有找到相应的key就默认初始化