使用Transformer进行英文书籍类型 (多) 分类_transformer能做多分类吗-CSDN博客

本文链接：https://blog.csdn.net/lohaive1015/article/details/121734064

本文介绍了一个使用Transformer模型进行英文书籍四分类（恐怖、科幻、幽默、犯罪）的任务。选择了预训练模型“bert-base-uncased”，由于资源限制。详细步骤包括数据读取、超参数搜索、模型训练和测试。最终测试的macro F1值达到0.72+，在众多参与者中排名40多位。

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

问题描述：对英文书籍进行分类，类型一共有四种：恐怖，科幻，幽默和犯罪。

训练数据集：英文书籍的一部分文字及对应标签，以及书籍的id，保存在工作路径下。

样本举例如下：

使用的预训练模型："bert-base-uncased"，为啥选这个模型嘞，因为它比较小，毕竟用Transformer太费算力了，自己租GPU压力还是很大的(Ｔ▽Ｔ)

1 import使用的packages

import torch
from transformers import BertTokenizerFast, BertForSequenceClassification, BertConfig
from transformers import Trainer, TrainingArguments
import numpy as np
import random
from sklearn.model_selection import train_test_split
import json
from sklearn.metrics import f1_score
from functools import reduce
torch.cuda.current_device()
torch.cuda._initialized = True

2 读取数据，这一次有单独的测试数据，因此只把训练数据集分成了训练集和验证集两个部分

# 读取训练数据集
train_data = json.load(open("traindata.json", "r")) #读取后为dictionary格式
X = train_data['X'] #书籍部分文本
Y = train_data['Y'] #分类标签，共4类，分别对应0,1,2,3
docid = train_data['docid'] #书籍的id，注意训练数据中书籍会重复出现

# 读取测试数据
test_data = json.load(open("testdata.json", "r"))
Xt = test_data['X']

# 将训练数据集随机分成训练集和验证集，比例为8:2，并基于标签进行分层抽样
(X_train, X_val, Y_train, Y_val) = train_test_split(X, Y, test_size = 0.2, stratify = Y)

# 定义即将要使用的预训练模型
model_name = "bert-base-uncased"
# 设置一个最大的文本长度
max_length = 512

# 载入预训练模型的分词器
tokenizer = BertTokenizerFast.from_pretrained(model_name, do_lower_case = True)

# 用分词器来做encoding，词数少于最大文本长度的话用0填充
train_encodings = tokenizer(X_train, truncation = True, padding = True, max_length = max_length)
valid_encodings = tokenizer(X_val, truncation = True, padding = True, max_length = max_length)

# 应用torch.utils.data.Dataset类来生成数据集
class generate_dataset(torch.utils.data.Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels            

    def __getitem__(self, index):