我们在 bert-base-chinese 中文预训练模型的基础上进行微调，以适应在新的数据集-中文酒店评论上进行文本分类。在这里我们使用 Train 类来完成中文评论分类模型的训练。我们训练时，只训练下游任务的参数部分。

from datasets import load_from_disk
from transformers import BertForSequenceClassification
from transformers import TrainingArguments
from transformers import Trainer
from sklearn.metrics import accuracy_score, f1_score, recall_score, precision_score
from sklearn.metrics import precision_recall_fscore_support

1. TrainingArguments

Trainer 要求我们传递一些训练相关的参数，即: 构建 TrainingArguments 对象，其具体参数如下：

1. output_dir 参数为必须添加的参数，它主要指定了模型训练过程中需要输出的内容，例如：保存 checkpoint 的路径
2. per_device_train_batch_size 参数设置 batch_size 大小，由于 Bert 模型很大，如果在 GPU 训练时出现 Cuda 内存分配失败，则应该相应调低该值，我的显卡为 RTX 2060，显存大小为 6G，在训练时就报错了，所以，我把这个参数设置为 2
3. overwrite_output_dir 参数表示覆盖输出目录中的内容
4. num_train_epochs 参数表示训练的 Epoch 数量
5. learning_rate 这个不用说了，设置学习率
6. eval_steps 参数表示每次训练多少步时，使用设置的验证集进行一次评估。我这里设置的是 10000，这是因为我的总的迭代次数是 19 万多，设置太小，评估过于频繁了，这个根据自己情况设置就好
7. save_steps 参数表示训练迭代多少次，也就是训练多少个 batch 保存一次模型
8. logging_dir 表示日志的输出目录
9. metric_for_best_model 因为我们在训练过程中会输出多个 checkpoint，这个参数用来指定评估最优模型的指标是什么，我们这里使用的时 precision，如果不指定的话，默认是根据 loss 值来评定最优模型

2. Trainer

Trainer 是 Transformers 封装的用于模型训练的高级 API，可以帮助我们快速搭建训练过程。训练机器配置为：centos7.8 + torch1.8.2 + cuda11 + 10代i5 + 16G固态内存 + RTX 2060 显卡，训练时间大概在 3 小时 30 分钟左右。

我们设置的训练参数如下：

1. 训练轮数：20
2. 每 2000 次迭代存储一次模型
3. 初始学习率设置为 2e-5
4. 优化方法使用的是 AdamW
5. 每次迭代送入模型的数据是 32 个
6. 每次迭代更新一次梯度
7. 模型存储目录为当前的 model 目录下

完整训练代码如下：

def train():

    # 1. 读取数据集
    data = load_from_disk('data/senti-dataset')
    # 将数据类型转换为张量
    data.set_format('pytorch', columns=['labels', 'input_ids', 'token_type_ids', 'attention_mask'])
    # 训练集
    train_data = data['train'].remove_columns(['review', '__index_level_0__'])
    # 验证集
    # valid_data = data['valid'].remove_columns(['review', '__index_level_0__'])

    # 2. 训练器
    train_model = BertForSequenceClassification.from_pretrained(
        'bert-base-chinese',
        num_labels=2)
    train_model.to('cuda')
    # 固定基础编码器的参数
    for param in train_model.base_model.parameters():
        param.requires_grad = False


    train_args = TrainingArguments(output_dir='model',
                                   overwrite_output_dir=True,
                                   save_steps=2000,
                                   per_device_train_batch_size=32,
                                   num_train_epochs=20,
                                   learning_rate=2e-5,
                                   evaluation_strategy='no',  # 训练期间不进行评估
                                   # eval_steps=1000,
                                   logging_strategy='epoch',
                                   logging_dir='logs',
                                   gradient_accumulation_steps=1,  # 每n个step更新一次参数
                                   metric_for_best_model='accuracy')

    trainer = Trainer(train_model,
                      train_args,
                      train_dataset=train_data)

    # 训练模型
    trainer.train()
    # 存储模型
    trainer.save_model('model/checkpoint-final')

训练过程输出如下：

******************* Running training *******************
  Num examples = 50216
  Num Epochs = 20
  Instantaneous batch size per device = 32
  Total train batch size (w. parallel, distributed & accumulation) = 32
  Gradient Accumulation steps = 1
  Total optimization steps = 31400
*******************************************************
{'loss': 0.5807, 'learning_rate': 1.9e-05, 'epoch': 1.0}
{'loss': 0.4634, 'learning_rate': 1.8e-05, 'epoch': 2.0}
{'loss': 0.4168, 'learning_rate': 1.7e-05, 'epoch': 3.0}
{'loss': 0.3925, 'learning_rate': 1.6e-05, 'epoch': 4.0}
{'loss': 0.3779, 'learning_rate': 1.5e-05, 'epoch': 5.0}
{'loss': 0.366,  'learning_rate': 1e-05,   'epoch': 6.0}
{'loss': 0.3604, 'learning_rate': 1e-05,   'epoch': 7.0}
{'loss': 0.355,  'learning_rate': 1e-05,   'epoch': 8.0}
{'loss': 0.3507, 'learning_rate': 1e-05,   'epoch': 9.0}
{'loss': 0.3486, 'learning_rate': 1e-05,   'epoch': 10.0}
{'loss': 0.344,  'learning_rate': 9e-06,   'epoch': 11.0}
{'loss': 0.3427, 'learning_rate': 8e-06,   'epoch': 12.0}
{'loss': 0.341,  'learning_rate': 7e-06,   'epoch': 13.0}
{'loss': 0.3393, 'learning_rate': 6e-06,   'epoch': 14.0}
{'loss': 0.3404, 'learning_rate': 5e-06,   'epoch': 15.0}
{'loss': 0.3395, 'learning_rate': 4e-06,   'epoch': 16.0}
{'loss': 0.3386, 'learning_rate': 3e-06,   'epoch': 17.0}
{'loss': 0.3367, 'learning_rate': 2e-06,   'epoch': 18.0}
{'loss': 0.3361, 'learning_rate': 1e-06,   'epoch': 19.0}
{'loss': 0.34,   'learning_rate': 0.0,     'epoch': 20.0}

3. 训练日志

另外，模型训练结束后会在 logging_dir 目录下输出 events.out.tfevents文件，该文件可以通过 tensorboard 可视化工具来查看输出日志，命令为：

tensorboard --logdir=./logs --port=8008