PyTorch ONNX & TensorRT

最新推荐文章于 2025-03-07 10:06:43 发布

Sycret

最新推荐文章于 2025-03-07 10:06:43 发布

阅读量783

点赞数

分类专栏：学习笔记

本文链接：https://blog.csdn.net/qq_34131692/article/details/111357975

版权

本文介绍了将PyTorch模型通过ONNX转换为TensorRT的过程，以实现模型量化和推理加速。首先，详细阐述了Model Quantization的一般流程，接着讨论了ONNX作为模型转换中介的作用和其支持的框架。接着展示了如何使用torch.onnx.export()将PyTorch模型导出为ONNX格式。最后，提到将ONNX模型转换为TensorRT时遇到的问题和解决思路，尽管因平台权限未能完成转换，但指出了后续步骤和可能需要的手动调整。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

Model Quantization

Model Quantization的一个比较经典的流程：对PyTorch训练好的model，先转onnx，再转成比如TensorRT（即quantization部分）。ONNX可以看成一个model的中转站，可以转换各种形式的model。（Quantization根据平台不同而不同，没有硬件的话自己做没有意义）

ONNX是一种针对机器学习所设计的开放式的文件格式，用于存储训练好的模型。它使得不同的人工智能框架（如Pytorch, MXNet）可以采用相同格式存储模型数据并交互。 ONNX的规范及代码主要由微软，亚马逊，Facebook 和 IBM 等公司共同开发，以开放源代码的方式托管在Github上。目前官方支持加载ONNX模型并进行推理的深度学习框架有： Caffe2, PyTorch, MXNet，ML.NET，TensorRT 和 Microsoft CNTK，并且 TensorFlow 也非官方的支持ONNX。—维基百科

那么首先写一个把PyTorch转换成ONNX的文件。

查了一下，要做转换的话好像这个model本身要走一遍，因此把train.py中初始化model以及之前用到的库都得包括进来，另外因为用了百度的平台，使用永久外部库的安装方式，所以最末还加了一行。

import torch
import torchvision.models as models

from utils.datasets import *
from utils.parse_config import *
from test import evaluate
from utils.utils import *
from models import *

import os
import sys
import time
import datetime
import argparse

sys.path.append('../../../aistudio/external-libraries')

这边还是和train.py中差不多一样的初始化model

if __name__ == "__main__":  
    parser = argparse.ArgumentParser()

    parser.add_argument("--batch_size", type=int, default

最低0.47元/天解锁文章