本文使用 https://huggingface.co/models 中的bert-base-uncased预训练模型和aclImdb数据集进行实战。

部分库版本#

transformers==4.17.0
torch==1.11.0
datasets==2.0.0

数据介绍#

AclImdb – v1 Dataset 是用于二进制情绪分类的大型电影评论数据集，其涵盖比基准数据集更多的数据，其中有 25,000 条电影评论用于训练，25,000 条用于测试，还有其他未经标记的数据可供使用。

本次只使用里面的训练数据和测试数据，使用的结构如下：

1
|- test
2
|-- neg
3
|-- pos
4
|- train
5
|-- neg
6
|-- pos

数据预处理#

要想使用Trainer进行训练，需要将数据调整到一定的规范，以下展示使用预训练模型对应的文本标记器（tokenizer）和datasets库处理原始数据。

导入必要库#

1
from datasets import Dataset
2
from transformers import BertTokenizer
3
import os

载入AclImdb数据，并转化为datasets.Dataset#

1
# 载入原始数据
2
def load_data(base_path):
3
    paths = os.listdir(base_path)
4
    result = []
5
    for path in paths:
6
        with open(os.path.join(base_path, path), 'r', encoding='utf-8') as f:
7
            result.append(f.readline())
8
    return result
9

10
# 读入数据并转化为datasets.Dataset
11
def get_dataset(base_path):
12
    # 为了展示方便，这里只取前3个数据，真实使用需要删掉切片操作
13
    pos_data = load_data(os.path.join(base_path, 'pos'))[:3]
14
    neg_data = load_data(os.path.join(base_path, 'neg'))[:3]
15

16
    # 列表合并
17
    texts = pos_data + neg_data
18
    # 生成标签，其中使用 '1.' 和 '0.' 是因为需要转化为浮点数，要不然模型训练时会报错
19
    labels = [[1., 0.]]*len(pos_data) + [[0., 1.]] * len(neg_data)
20
    dataset = Dataset.from_dict({'texts':texts, 'labels':labels})
21
    return dataset
22

23
# 加载数据
24
train_dataset = get_dataset('../data/aclImdb/train/')
25
test_dataset = get_dataset('../data/aclImdb/test/')

此时输出train_dataset可以查看其结构：

1
print(train_dataset)
2

3
输出结果
4
Dataset({
5
    features: ['texts', 'labels'],
6
    num_rows: 6
7
})

可以看到，包含了texts和labels，有6行数据。

通过字典的方式可以查看其中的元素：

1
print(train_dataset['labels'])
2

3
输出结果：
4
[[1.0, 0.0], [1.0, 0.0], [1.0, 0.0], [0.0, 1.0], [0.0, 1.0], [0.0, 1.0]]

通过.features可以查看特征：

1
print(train_dataset.features)
2

3
输出结果：
4
{
5
'texts': Value(dtype='string', id=None),
6
'labels': Sequence(feature=Value(dtype='float64', id=None), length=-1, id=None)
7
}

载入文本标记器#

1
# cache_dir是预训练模型的地址
2
cache_dir="../transformersModels/bert-base-uncased/"
3
tokenizer = BertTokenizer.from_pretrained(cache_dir)

注意

这个路径的模型要自己下载，不能是transformer包下的，要不会报错。

将数据转化为模型可以接受的格式#

1
# 设置最大长度
2
MAX_LENGTH = 512
3

4
# 使用文本标记器对texts进行编码
5
train_dataset = train_dataset.map(lambda e: tokenizer(e['texts'], truncation=True, padding='max_length', max_length=MAX_LENGTH), batched=True)
6
test_dataset = test_dataset.map(lambda e: tokenizer(e['texts'], truncation=True, padding='max_length', max_length=MAX_LENGTH), batched=True)

运行完后，数据集会包含训练模型所需要的所有参数：

1
print(print(train_dataset.features))
2

3
输出结果：
4
{
5
'texts': Value(dtype='string', id=None),
6
 'labels': Sequence(feature=Value(dtype='float64', id=None), length=-1, id=None),
7
 'input_ids': Sequence(feature=Value(dtype='int32', id=None), length=-1, id=None),
8
 'token_type_ids': Sequence(feature=Value(dtype='int8', id=None), length=-1, id=None),
9
 'attention_mask': Sequence(feature=Value(dtype='int8', id=None), length=-1, id=None)
10
}

保存处理好的数据到本地#

在数据量大的时候，处理数据需要很长的时间，为了不每次都重新处理数据，可以将数据先存到本地

1
train_dataset.save_to_disk('./data/train_dataset')
2
test_dataset.save_to_disk('./data/test_dataset')

训练模型#

导入必要库#

1
from transformers import BertForSequenceClassification, BertTokenizer, Trainer, TrainingArguments, BertConfig
2
import torch
3
from datasets import Dataset
4
import json
5
import os
6
# 设定使用的GPU编号，也可以不设置，但trainer会默认使用多GPU
7
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

载入网络并设置自定义尾部#

1
# 将num_labels设置为2，因为我们训练的任务为2分类
2
model = BertForSequenceClassification.from_pretrained('../transformersModels/bert-base-uncased/', num_labels=2)

加载处理好的数据#

1
train_dataset = Dataset.load_from_disk('./data/train_dataset/')
2
test_dataset = Dataset.load_from_disk('./data/test_dataset/')

冻结BERT参数#

1
for param in model.base_model.parameters():
2
    param.requires_grad = False

因为BERT是预训练模型，因此可以不再进行权重更新，只对尾部的分类器进行优化。与此同时，这个设置也会减少训练时使用的时间和显存。

创建trainer#

1
# 训练超参配置
2
training_args = TrainingArguments(
3
    output_dir='./my_results',          # output directory 结果输出地址
4
    num_train_epochs=10,              # total # of training epochs 训练总批次
5
    per_device_train_batch_size=128,  # batch size per device during training 训练批大小
6
    per_device_eval_batch_size=128,   # batch size for evaluation 评估批大小
7
    logging_dir='./my_logs',            # directory for storing logs 日志存储位置
8
)
9

10
# 创建Trainer
11
trainer = Trainer(
12
    model=model.to('cuda'),              # the instantiated 🤗 Transformers model to be trained 需要训练的模型
13
    args=training_args,                  # training arguments, defined above 训练参数
14
    train_dataset=train_dataset,         # training dataset 训练集
15
    eval_dataset=test_dataset,           # evaluation dataset 测试集
16
)

训练、评估和保存模型#

1
# 开始训练
2
trainer.train()
3

4
# 开始评估模型
5
trainer.evaluate()
6

7
# 保存模型 会保存到配置的output_dir处
8
trainer.save_model()

保存模型会生成三个文件：

1
# 模型配置文件
2
config.json
3

4
# 模型数据文件
5
pytorch-model.bin
6

7
# 训练配置文件
8
training_args.bin

查看日志#

使用tensorboard查看配置的logging_dir即可看到训练时的日志。

加载模型#

1
output_config_file = './my_results/config.json'
2
output_model_file = './my_results/pytorch_model.bin'
3

4
config = BertConfig.from_json_file(output_config_file)
5
model = BertForSequenceClassification(config)
6
state_dict = torch.load(output_model_file)
7
model.load_state_dict(state_dict)

载入成功后会显示：

1
<All keys matched successfully>

这个model就可以跟普通的pytorch模型一样使用了

如下：

1
cache_dir="../transformersModels/bert-base-uncased/"
2
tokenizer = BertTokenizer.from_pretrained(cache_dir)
3
data = tokenizer(['This is a good movie', 'This is a bad movie'], max_length=512, truncation=True, padding='max_length', return_tensors="pt")
4
print(model(**data))
5

6
输出结果：
7
SequenceClassifierOutput(
8
loss=None, logits=tensor([
9
[-0.2951,  0.5463],
10
[-0.4638,  0.6353]],
11
grad_fn=<AddmmBackward0>),
12
hidden_states=None,
13
attentions=None)

由于只用3条数据训练了10轮，因此结果很差，正常训练结果可以变好了。

完整代码#

处理数据#

1
from datasets import Dataset
2
from transformers import BertTokenizer
3
import os
4

5
def load_data(base_path):
6
    paths = os.listdir(base_path)
7
    result = []
8
    for path in paths:
9
        with open(os.path.join(base_path, path), 'r', encoding='utf-8') as f:
10
            result.append(f.readline())
11
    return result
12

13
def get_dataset(base_path):
14
    pos_data = load_data(os.path.join(base_path, 'pos'))
15
    neg_data = load_data(os.path.join(base_path, 'neg'))
16

17
    texts = pos_data + neg_data
18
    labels = [[1., 0.]]*len(pos_data) + [[0., 1.]] * len(neg_data)
19
    dataset = Dataset.from_dict({'texts':texts, 'labels':labels})
20
    return dataset
21

22
# 加载数据
23

24
train_dataset = get_dataset('../data/aclImdb/train/')
25
test_dataset = get_dataset('../data/aclImdb/test/')
26

27
# 处理数据
28
cache_dir="../transformersModels/bert-base-uncased/"
29
tokenizer = BertTokenizer.from_pretrained(cache_dir)
30

31
MAX_LENGTH = 512
32
train_dataset = train_dataset.map(lambda e: tokenizer(e['texts'], truncation=True, padding='max_length', max_length=MAX_LENGTH), batched=True)
33
test_dataset = test_dataset.map(lambda e: tokenizer(e['texts'], truncation=True, padding='max_length', max_length=MAX_LENGTH), batched=True)
34

35
# 保存数据
36
train_dataset.save_to_disk('./data/train_dataset')
37
test_dataset.save_to_disk('./data/test_dataset')

训练模型#

1
from transformers import BertForSequenceClassification, BertTokenizer, Trainer, TrainingArguments, BertConfig
2
import torch
3
from datasets import Dataset
4
import json
5
import os
6
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
7

8
# 加载模型
9
model = BertForSequenceClassification.from_pretrained('../transformersModels/bert-base-uncased/', num_labels=2)
10

11
# 加载数据
12
train_dataset = Dataset.load_from_disk('./data/train_dataset/')
13
test_dataset = Dataset.load_from_disk('./data/test_dataset/')
14

15
# 冻结bert参数
16
for param in model.base_model.parameters():
17
    param.requires_grad = False
18

19
# 创建trianer
20
training_args = TrainingArguments(
21
    output_dir='./my_results',          # output directory 结果输出地址
22
    num_train_epochs=10,              # total # of training epochs 训练总批次
23
    per_device_train_batch_size=128,  # batch size per device during training 训练批大小
24
    per_device_eval_batch_size=128,   # batch size for evaluation 评估批大小
25
    logging_dir='./my_logs',            # directory for storing logs 日志存储位置
26
)
27

28
trainer = Trainer(
29
    model=model.to('cuda'),              # the instantiated 🤗 Transformers model to be trained 需要训练的模型
30
    args=training_args,                  # training arguments, defined above 训练参数
31
    train_dataset=train_dataset,         # training dataset 训练集
32
    eval_dataset=test_dataset,           # evaluation dataset 测试集
33
)
34

35
# 训练和保存
36
trainer.train()
37
trainer.save_model()

载入模型#

1
output_config_file = './my_results/config.json'
2
output_model_file = './my_results/pytorch_model.bin'
3

4
config = BertConfig.from_json_file(output_config_file)
5
model = BertForSequenceClassification(config)
6
state_dict = torch.load(output_model_file)
7
model.load_state_dict(state_dict)

原文链接：CSDN