首页 > 其他分享 >pytorch 转 tensorRT 踩的几个小坑_tensorrt engine set up failed

pytorch 转 tensorRT 踩的几个小坑_tensorrt engine set up failed

时间:2024-05-19 12:58:53浏览次数:27  
标签:engine set target onnx torch up output data size

CSDN搬家失败,手动导出markdown后再导入博客园

1、版本不匹配

[E] [TRT] Layer:Where_51's output can not be used as shape tensor.
[E] [TRT] Network validation failed.
[E] Engine creation failed.
[E] Engine set up failed.

这实际是由于 pytorch 与 TensorRT 版本不匹配,我的 TensorRT 是 7.0,pytorch 应该是 1.4,但我用了 1.7

因此需要用 1.7 重新读取权重文件,然后用老的方式保存,再用 onnx 导出

def main():
    input_shape = (3, 416, 416)
    model_onnx_path = "yolov4tiny.onnx"
 
    # model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch', 'unet',
    #                        in_channels=3, out_channels=1, init_features=32, pretrained=True)
    model = YoloBody(3, 12).cuda()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    dummy_input = torch.randn(1, 3, 416, 416, device=device)
    # 用1.7版本读取权重
    state_dict = torch.load('logs/Epoch120-Total_Loss0.5324-Val_Loss0.8735.pth', map_location=device)
    model.load_state_dict(state_dict)
    # 保存成1.4版本支持的格式
    torch.save(model.state_dict(), 'logs/for_onnx.pth', _use_new_zipfile_serialization=False)
    
    # Python解释器换成torch1.4的环境,重新读取,导出pytorch1.4版本对应的onnx
    state_dict = torch.load('logs/for_onnx.pth', map_location=device)
    model.load_state_dict(state_dict)
    model.train(False)
 
    inputs = ['input_1']
    outputs = ['output_1', 'output_2']
    dynamic_axes = {'input_1': {0: 'batch'}, 'output_1': {0: 'batch'}}
    torch.onnx.export(model,
                      dummy_input,
                      model_onnx_path,
                      export_params=True,
                      opset_version=11,
                      do_constant_folding=True,
                      input_names=inputs, output_names=outputs,
                      dynamic_axes=None)

这样操作产生的 onnx 文件才能被 TensorRT 转换

使用 TensorRT 的 OSS 工具:

class Upsample(nn.Module):
    def __init__(self):
        super(Upsample, self).__init__()
 
    def forward(self, x, target_size, inference=False):
        assert (x.data.dim() == 4)
        # _, _, tH, tW = target_size
 
        if inference:
 
            #B = x.data.size(0)
            #C = x.data.size(1)
            #H = x.data.size(2)
            #W = x.data.size(3)
 
            return x.view(x.size(0), x.size(1), x.size(2), 1, x.size(3), 1).\
                    expand(x.size(0), x.size(1), x.size(2), target_size[2] // x.size(2), x.size(3), target_size[3] // x.size(3)).\
                    contiguous().view(x.size(0), x.size(1), target_size[2], target_size[3])
        else:
            return F.interpolate(x, size=(target_size[2], target_size[3]), mode='nearest')

建议打开 --verbose,转换过程会很慢,verbose 打印日志看着能安心点,要不然盯着屏幕会以为卡死了,发慌

2、上采样 scale 问题

[5] Assertion failed: ctx->tensors().count(inputName)

![[output/attachments/d697b05ccbc0cd8ef1baff4f153b9024_MD5.png]]

YOLO 的上采样阶段,Pytorch 使用 opset=11 的 onnx 会导致 upsample 层里增加一个 constant 节点,所以 TensorFlowRT 转换失败,期间参考 pytorch 经 onnx 使用 TensorRT 部署转换踩坑记录中提到的方法,无效

![[output/attachments/95e6ae1c7b9192edb8b6674198bb354e_MD5.png]]

尝试多个版本的 Pytorch 与 onnx 后,upsample 层的问题依然解决不了,最后参考 https://github.com/Tianxiaomo/pytorch-YOLOv4 这个实现,在 inference 时不使用 torch 自己的插值函数,而是自己重写,成功导出 TensorRT

class Upsample(nn.Module):
    def __init__(self):
        super(Upsample, self).__init__()
 
    def forward(self, x, target_size, inference=False):
        assert (x.data.dim() == 4)
        # _, _, tH, tW = target_size
 
        if inference:
 
            #B = x.data.size(0)
            #C = x.data.size(1)
            #H = x.data.size(2)
            #W = x.data.size(3)
 
            return x.view(x.size(0), x.size(1), x.size(2), 1, x.size(3), 1).\
                    expand(x.size(0), x.size(1), x.size(2), target_size[2] // x.size(2), x.size(3), target_size[3] // x.size(3)).\
                    contiguous().view(x.size(0), x.size(1), target_size[2], target_size[3])
        else:
            return F.interpolate(x, size=(target_size[2], target_size[3]), mode='nearest')

3、数据类型错误

Unsupported ONNX data type: DOUBLE (2)

由于上面的改动中,出现了除法,导致 TRT 识别双精度,下图中出现的 cast 节点即为问题。很神奇的是用上面原始的代码不会出现这个问题,换了我自己的模型就有问题了

![[output/attachments/20922cff4c9b9fce378baf949d42b34f_MD5.png]]

根据 https://github.com/onnx/onnx-tensorrt/issues/400#issuecomment-730240546 的说法

![[output/attachments/960f5b584c8b980e4c9713b43a46a8a4_MD5.png]]

尝试了下并没什么用。

为了解决这个问题,直接把 target_size[3] // x.size(3) 换成结果 2 即可,成功。

标签:engine,set,target,onnx,torch,up,output,data,size
From: https://www.cnblogs.com/algorithmSpace/p/18200236

相关文章

  • [MASM拾遗]Offset
      Offset伪指令我一直都认为只是获取标识符在段中的偏移地址,但经研究,发现了部分违反直觉的细微区别:  1、在完整端声明(Fullsegmentdefinition)的情况下,如果offsetmygroup:myvar或offsetmysegment:myvar,可通过端前缀来获取myvar与group开头位置的偏移地址或myvar与mysegme......
  • BeautifulSoup库
    一、安装BeautifulSoup库 可以现在目前python安装了哪些包安装beautifulsoup二、beautifulsoup官网https://www.crummy.com/software/BeautifulSoup/bs4/doc/三、beautifulsoup的主要解析器 四、beautifulsoup的find函数查找html的titlefrombs4importBeautifulS......
  • puppeteer使用一些技巧简单说明
    puppeteer是一个nodejs包提供了方便的基于devtools协议进行chrome/chromium控制,puppeteer默认运行在无头模式以下是对于puppeteer使用的一些简单总结一些问题browser&&page对象复用问题实际上还是结合实际,个人建议减少复用,除非自己对于browser&&page进行了比较......
  • puppeteer localstorage 处理简单说明
    puppeteer对于cookie以及localstorage的控制实际上还是比较常见的需求,现在主流web开发都比较喜欢使用localstorage进行一些数据的额存储,以下是关于localstorage处理的一个简单演示环境准备基于了browserless(使用了v2版本,注意browserlessv1与v2是有差异的,目前测试chrome......
  • Vue3报错:已声明“upperName”,但从未读取其值。ts-plugin(6133)
    Vue3报错:已声明“upperName”,但从未读取其值。ts-plugin(6133)报错显示:类型“StoreToRefs<Store<"count",{sum:number;name:string;address:string;},{},{increment(value:number):void;}>>”上不存在属性“upperName”。ts-plugin(2339)相关代码:vue文件:con......
  • C# app.config配置appSettings标签
    AppSettings标签主要用于应用程中的一些配置信息。比如上传文件的路径,需要用到的可能会根据不同环境改变的一些常量等。可以在配置文件添加可执行文件要的配置。usingSystem;usingSystem.Collections.Generic;usingSystem.Configuration;classProgram{staticvoid......
  • VUE速通(10)Vue3核心语法(2)setup
    1setup概述setup是Vue3中一个新的配置项,值是一个函数,它是CompositionAPI“表演的舞台”,组件中所用到的:数据、方法、计算属性、监视......等等,均配置在setup中。特点如下:setup函数返回的对象中的内容,可直接在模板中使用。setup中访问this是undefined。setup函数会......
  • bitset专题
    bitsetbitset前身:普通状态压缩的优化以cf937G为例,对于邻接矩阵的由二维压缩到一维#include<bits/stdc++.h>usingi64=longlong;voidsolve(){intn;std::cin>>n;std::vector<std::string>g(n),w(n);for(inti=0;i<n;i++){......
  • [lnsyoj281/luoguP2023/AHOI2009]维护序列
    题意原题链接给定序列\(a\),要求维护区间加,区间乘,区间查询三种操作sol显然线段树,事实上,这是一道板子题(luoguP3373),但由于蒟蒻实在是太蒻了,并没有打过这道题。区间加如果我们将区间里的每一个元素都插入线段树做一次修改操作,那么一次修改操作的时间复杂度为\(O(n\logn)\),此时......
  • ctflearn-writeup(Exclusive Santa)
    https://ctflearn.com/challenge/851在完成这题前最好先下载foremost,unrar,stegsolve等工具首先拿到题目后,先解压得到两个图片文件1.png和3.png两张图片用exiftool,strings,binwalk试过后发现无解于是用关键命令foremost3.png-T(修复破损文件)发现有一个output的文件夹......