Transformers文本分类微调和TensorRT推理加速

Transformers文本分类微调和TensorRT推理加速

本文介绍了使用Transformers模块微调文本分类模型，以及模型转pth格式，进而转为onnx格式和TensorRT的engine格式，最后通过TensorRT进行推理的完整流程

开发环境简介

宿主机环境

Linux version:  Ubuntu20.04
Docker version: 20.10.20
Python version: 3.8.5 (3.6-3.9)
NVIDIA version: 1080Ti
Cuda version: cuda11.6 (>11.1)
Driver version: 510.47.03 (>460)

Python packages:

transformers           4.24.0
torch                  1.13.0+cu116
scikit-learn           1.1.3
pandas                 1.5.1
numpy                  1.23.5

容器内环境

容器由TensorRT DockerFile构建，将在后续进行介绍，这里只介绍python环境

Python version