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

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

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

PS: 大模型基础和进阶付费课程（自己讲的）：《AIGC大模型理论与工业落地实战》-CSDN学院 或者《AIGC大模型理论与工业落地实战》-网易云课堂

开发环境简介

宿主机环境

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: 3.6.9

Python packages:

numpy               1.19.5
pycuda              2022.1
torch               1.10.2+cu111
transformers        4.18.0

一. 准备数据

1.数据样例

数据集可以从 数据样例这里进行下载

(base) ys@ys:~/Documents/my_trt/data$ head -10 dev.csv 
label,txt
-1,一个月都有点卡了，
1,手机很不错，玩游戏很流畅，快递小哥态度特别特别特别好，重要的事情说三遍?
1,初步用了一下，开始拿到手觉得真**大，玩了一会以后 真香！玩游戏很爽，双扬声器看片也很震撼～拍照的话现在还凑活看后期更新优化吧！这款手机比较适合影音娱乐用，一般人不建议用。
-1,续航不行 玩吃鸡卡
1,物流配送速度很快，三天左右就到了，画面清晰，音质还挺好的，总之超喜欢的
0,着好看，都挺不错的。   但是京东的保价，双十一便宜100元，联系客服，说退100，这都10多天了，仍然没有解决。
1,手机是正品，质量很好，速度很快，价格不比别家便宜，还送膜、送耳机、送自拍神器。谢谢商家！
1,机很流畅，买的很好的一次，下次还要购买，一点都不卡。
0,不能再相信它了

标注为1表示情感积极，标注为0表示情感中立，标注为-1表示情感消极

(base) ys@ys:~/Documents/my_trt/data$ tree
.
├── dev.csv
├── test.csv
└── train.csv

2.数据处理

#!/usr/bin/env python
# coding: utf-8

# ## 读取本地数据

import pandas as pd

df_train = pd.read_csv("data/train.csv")
df_dev   = pd.read_csv("data/dev.csv")
df_test  = pd.read_csv("data/test.csv")

dict_key = {
   -1:2, 0:0, 1:1}  # 标签为-1的改为2，其他保持不变；标签-1在Transformers中读取会报错
df_train["label"] = df_train["label"].apply(lambda x: dict_key[x])
df_dev["label"]   = df_dev["label"].apply(lambda x: dict_key[x])
df_test["label"]  = df_test["label"].apply(lambda x: dict_key[x])

3.构造DataSets格式数据

import os
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
from transformers import Trainer, TrainingArguments
from transformers import pipeline, DataCollatorWithPadding
from datasets import load_dataset, Dataset, Features, Value, ClassLabel

# 构造dataset格式数据集

class_names = [0, 1, 2]
features = Features({
   'txt': Value('string'), 'label': ClassLabel(names=class_names)})

train_dataset = Dataset.from_pandas(df_train, split="train", features=features)
dev_dataset = Dataset.from_pandas(df_test, split="dev", features=features)
test_dataset = Dataset.from_pandas(df_test, split="test", features=features)

二.下载和加载预训练模型

1.下载bert预训练模型参数

前往https://huggingface.co/models，搜索bert-base-chinese，在Files and versions中下载三个文件

(base) ys@ys:~/Documents/my_trt/bert-base-chinese$ tree
.
├── config.json
├── pytorch_model.bin
└── vocab.txt

2.加载bert预训练模型参数