weblog

Appearance

预训练BERT ​

一、概述 ​

之前我们已经学了 BERT 的模型结构，还有用来预训练 BERT 的 WikiText-2 数据集，这一节我们就来实际动手预训练 BERT，并且看看怎么用预训练好的 BERT 来表示文本。

原始的 BERT 有两个版本：BERT BASE 和 BERT LARGE，BERT BASE 有 12 层、768 个隐藏单元、12 个自注意力头，有 1.1 亿个参数；BERT LARGE 有 24 层、1024 个隐藏单元、16 个自注意力头，有 3.4 亿个参数，参数非常多，训练起来需要很大的算力。为了方便演示，我们这里定义一个小的 BERT 模型，只用 2 层、128 个隐藏单元和 2 个自注意力头，这样训练起来更快，适合学习和测试。

二、定义小 BERT 模型 ​

我们用 d2l 库的 BERTModel 来定义这个小 BERT 模型，代码如下（以 PyTorch 为例）：

python
import torch
from torch import nn
from d2l import torch as d2l

# 假设我们已经有了词表vocab
net = d2l.BERTModel(
    len(vocab),  # 词表大小
    num_hiddens=128,  # 隐藏单元数
    norm_shape=[128],  # 层归一化的形状
    ffn_num_input=128,  # 前馈网络的输入维度
    ffn_num_hiddens=256,  # 前馈网络的隐藏层维度
    num_heads=2,  # 自注意力头数
    num_layers=2,  # Transformer编码器的层数
    dropout=0.2,  # dropout比例
    key_size=128,  # 注意力键的维度
    query_size=128,  # 注意力查询的维度
    value_size=128,  # 注意力值的维度
    hid_in_features=128,  # 隐藏层的输入维度
    mlm_in_features=128,  # 遮蔽语言模型任务的输入维度
    nsp_in_features=128  # 下一句预测任务的输入维度
)
# 获取所有可用的GPU
devices = d2l.try_all_gpus()
# 把模型放到GPU上
net = nn.DataParallel(net, device_ids=devices).to(devices[0])
# 定义损失函数，用交叉熵损失
loss = nn.CrossEntropyLoss()

三、预训练的损失计算 ​

BERT 的预训练有两个任务，所以损失是这两个任务的损失之和：

  1. 遮蔽语言模型（MLM）损失：就是随机遮住文本里的一些词，让模型预测这些被遮住的词，损失就是模型预测的结果和真实词的交叉熵损失。

  2. 下一句预测（NSP）损失：判断两个句子是不是连续的，损失是模型预测的结果和真实标签的交叉熵损失。

我们定义一个辅助函数_get_batch_loss_bert来计算这两个损失：

python
def _get_batch_loss_bert(net, loss, vocab_size, tokens_X, segments_X, valid_lens_x, pred_positions_X, mlm_weights_X, mlm_Y, nsp_y):
    # 前向传播，得到模型的输出
    # encoded_X是所有词元的表示，mlm_Y_hat是遮蔽语言模型的预测结果，nsp_Y_hat是下一句预测的结果
    encoded_X, mlm_Y_hat, nsp_Y_hat = net(tokens_X, segments_X, valid_lens_x.reshape(-1), pred_positions_X)
    
    # 计算遮蔽语言模型的损失
    # 把预测结果和真实标签都展平，然后计算交叉熵损失，还要乘以权重（因为有些位置是不需要预测的，权重为0）
    mlm_l = loss(mlm_Y_hat.reshape((-1, vocab_size)), mlm_Y.reshape(-1)) * mlm_weights_X.reshape((-1, 1))
    # 对损失求和，然后除以有效权重的和（避免除以0，加一个很小的数）
    mlm_l = mlm_l.sum() / (mlm_weights_X.sum() + 1e-8)
    
    # 计算下一句预测的损失，取均值
    nsp_l = loss(nsp_Y_hat, nsp_y).mean()
    
    # 总损失是两个损失的和
    total_l = mlm_l + nsp_l
    return mlm_l, nsp_l, total_l

四、训练 BERT ​

我们定义一个训练函数train_bert，用来训练 BERT 模型：

python
def train_bert(train_iter, net, loss, vocab_size, devices, num_steps):
    # 定义Adam优化器，学习率设为0.01
    trainer = torch.optim.Adam(net.parameters(), lr=0.01)
    step, timer = 0, d2l.Timer()
    # 用来绘制损失变化的动画
    animator = d2l.Animator(xlabel='step', ylabel='loss', xlim=[1, num_steps], legend=['mlm', 'nsp'])
    # 用来记录损失的累加器：mlm损失总和、nsp损失总和、处理的句子对数量、计数
    metric = d2l.Accumulator(4)
    num_steps_reached = False
    
    while step < num_steps and not num_steps_reached:
        # 遍历训练数据集
        for tokens_X, segments_X, valid_lens_x, pred_positions_X, mlm_weights_X, mlm_Y, nsp_y in train_iter:
            # 清空梯度
            trainer.zero_grad()
            timer.start()
            # 计算损失
            mlm_l, nsp_l, total_l = _get_batch_loss_bert(
                net, loss, vocab_size, tokens_X, segments_X, valid_lens_x,
                pred_positions_X, mlm_weights_X, mlm_Y, nsp_y
            )
            # 反向传播
            total_l.backward()
            # 更新参数
            trainer.step()
            # 记录损失和处理的句子对数量
            metric.add(mlm_l, nsp_l, tokens_X.shape[0], 1)
            timer.stop()
            # 绘制损失
            animator.add(step + 1, (metric[0] / metric[3], metric[1] / metric[3]))
            step += 1
            if step == num_steps:
                num_steps_reached = True
                break
    
    # 输出最终的损失
    print(f'MLM损失 {metric[0] / metric[3]:.3f}, NSP损失 {metric[1] / metric[3]:.3f}')
    # 输出训练速度
    print(f'每秒处理 {metric[2] / timer.sum():.1f} 个句子对，使用设备 {str(devices)}')

调用这个函数训练，比如训练 1000 步：

python
# 假设我们已经加载了训练数据集train_iter
train_bert(train_iter, net, loss, len(vocab), devices, 1000)

训练完成后，会输出 MLM 损失和 NSP 损失，还有训练的速度。

五、用 BERT 表示文本 ​

预训练好 BERT 之后，我们可以用它来表示单个文本或者文本对，得到每个词元的上下文敏感的表示。我们定义一个函数来获取文本的 BERT 表示：

python
def get_bert_encoding(net, tokens_a, tokens_b=None):
    # 获取文本的词元和片段索引
    tokens, segments = d2l.get_tokens_and_segments(tokens_a, tokens_b)
    # 把词元转换成词元id，放到GPU上，增加一个批量维度
    token_ids = torch.tensor(vocab[tokens], device=devices[0]).unsqueeze(0)
    # 把片段索引转换成张量，放到GPU上，增加一个批量维度
    segments = torch.tensor(segments, device=devices[0]).unsqueeze(0)
    # 有效长度，就是词元的数量
    valid_len = torch.tensor(len(tokens), device=devices[0]).unsqueeze(0)
    # 前向传播，得到编码后的表示
    encoded_X, _, _ = net(token_ids, segments, valid_len)
    return encoded_X

1. 表示单个文本 ​

python
# 文本：['a', 'little', 'cat', 'is', 'playing']
tokens_a = ['a', 'little', 'cat', 'is', 'playing']
# 获取BERT表示
encoded_a = get_bert_encoding(net, tokens_a)
# 输出第一个词元（<cls>）的表示，形状是(1, 1, 128)，1是批量大小，1是词元位置，128是隐藏单元数
print(encoded_a[:, 0, :])

2. 表示文本对 ​

python
# 第一个文本：['a', 'little', 'cat', 'is', 'playing']
# 第二个文本：['it', 'is', 'running', 'fast']
tokens_b = ['it', 'is', 'running', 'fast']
encoded_ab = get_bert_encoding(net, tokens_a, tokens_b)
# 输出第一个词元（<cls>）的表示
print(encoded_ab[:, 0, :])

我们会发现，同一个词在不同的上下文里的表示是不一样的，比如 “is” 在第一个文本和文本对里的表示不同，这说明 BERT 的表示是上下文敏感的，能根据不同的上下文生成不同的表示。

六、小结 ​

  1. 原始 BERT 有两个版本，BERT BASE 和 BERT LARGE，参数很多，训练需要大量算力，我们这里用了一个小的 BERT 模型来演示预训练过程。

  2. BERT 的预训练损失是遮蔽语言模型损失和下一句预测损失的和，这两个任务分别让模型学习词的语义和句子之间的关系。

  3. 预训练好的 BERT 可以用来表示单个文本、文本对或者其中的词元，而且表示是上下文敏感的，同一个词在不同上下文里的表示不同。

  4. 训练 BERT 的时候用 Adam 优化器，多 GPU 训练可以提高训练速度。

（注：文档部分内容可能由 AI 生成） 源地址

京ICP备2024093538号-1

自用笔记本