Bert预训练代码

预训练Bert

1.加载WikiText-2数据集作为小批量的预训练样本，用于遮蔽语言模型和下一句预测。批量大小是512，BERT输入序列的最大长度是64。注意，在原始BERT模型中，最大长度是512。

2.原始BERT :cite:Devlin.Chang.Lee.ea.2018有两个不同模型尺寸的版本。基本模型使用12层（Transformer编码器块），768个隐藏单元（隐藏大小）和12个自注意头。大模型使用24层，1024个隐藏单元和16个自注意头。值得注意的是，前者有1.1亿个参数，后者有3.4亿个参数。为了便于演示，我们定义了一个小的BERT，使用了2层、128个隐藏单元和2个自注意头。

3.在定义训练代码实现之前，我们定义了一个辅助函数_get_batch_loss_bert。给定训练样本，该函数计算遮蔽语言模型和下一句子预测任务的损失。请注意，BERT预训练的最终损失是遮蔽语言模型损失和下一句预测损失的和。

4.下面的train_bert函数定义了在WikiText-2（train_iter）数据集上预训练BERT（net）的过程。训练BERT可能需要很长时间。以下函数的输入num_steps指定了训练的迭代步数，而不是像train_ch13函数那样指定训练的轮数（参见 :numref:sec_image_augmentation）。

python 复制代码

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

# ==========================================
# 1. 初始化超参数与微型模型
# ==========================================
# 这里的 batch_size 是 512，意味着每次同时给模型看 512 对句子。
# max_len 是 64，表示每对句子被截断或补齐到 64 个词元长度。
batch_size, max_len = 512, 64
train_iter, vocab = d2l.load_data_wiki(batch_size, max_len)  

# 实例化一个"微型版"的 BERT 模型用于演示（为了能在普通电脑或单卡上跑起来）
# 标准的 BERT-Base hidden_size 是 768，这里缩小到了 128；层数从 12 降到了 2 层。
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, dropout=0.2, key_size=128, query_size=128,
                    value_size=128, hid_in_features=128, mlm_in_features=128,
                    nsp_in_features=128)

devices = d2l.try_all_gpus() # 获取所有可用的显卡 (GPU)

"""分类任务"""
# 定义交叉熵损失函数 (用于算选择题的得分)
# 注意：如果要让下面的 weights 乘法生效，通常这里底层应该带有 reduction='none' 的属性，
# 从而保留每个预测词的独立误差。
loss = nn.CrossEntropyLoss()  

# ==========================================
# 2. 定义批次损失计算函数 (判卷老师)
# ==========================================
#@save
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):

"""
BERTModel 终极组装（Canvas 里的 bert_model.py）的 forward 函数最后是怎么返回的？
它返回了三个东西：
return encoded_X, mlm_Y_hat, nsp_Y_hat

encoded_X: 经过 Transformer 编码后的底层特征矩阵。
mlm_Y_hat: 完形填空的预测结果。
nsp_Y_hat: 下一句预测的预测结果。
"""
    # 步骤 A：前向传播 (让学生做题)
    # 把题目发下去，得到 MLM(完形填空) 和 NSP(下一句判断) 的答案
    _, mlm_Y_hat, nsp_Y_hat = net(tokens_X, segments_X,
                                  valid_lens_x.reshape(-1),
                                  pred_positions_X)
                                  
"""
这里在计算mlm_l时应该指定reduction='none'，否则计算出来是一个标量，weight就没起到该有的作用了
"""
    # 步骤 B：计算 MLM 完形填空的误差 (批改完形填空)
    # 先把预测结果和真实答案都压扁成一条长线，方便一一对比
    mlm_l = loss(mlm_Y_hat.reshape(-1, vocab_size), mlm_Y.reshape(-1))
    
    # 【核心妙招：过滤凑数坑位】
    # 还记得之前用 0 补齐的那些 MLM 位置吗？它们在 mlm_weights_X 里的权重是 0！
    # 这一步把误差乘以权重，瞬间就把那些"凑数坑位"产生的误差清零了，只保留真正挖掉的词的误差。
    mlm_l = mlm_l * mlm_weights_X.reshape(-1, 1)
    
    # 算出所有真实挖空词的平均误差 (加 1e-8 是为了防止分母为 0 导致报错)
    mlm_l = mlm_l.sum() / (mlm_weights_X.sum() + 1e-8)
    
    # 步骤 C：计算 NSP 下一句预测的误差 (批改判断题)
    nsp_l = loss(nsp_Y_hat, nsp_y)
    
    # 步骤 D：总成绩 = 完形填空误差 + 判断题误差
    # BERT 的强大就在于：同时被两个老师(任务)毒打，最终学到既懂词义又懂逻辑的语言表示！
    l = mlm_l + nsp_l
    return mlm_l, nsp_l, l  


# ==========================================
# 3. 核心训练主循环 (让机器疯狂刷题进化)
# ==========================================
def train_bert(train_iter, net, loss, vocab_size, devices, num_steps):
    # 将模型复制到多张显卡上并行加速 (如果有的话)
    net = nn.DataParallel(net, device_ids=devices).to(devices[0])
"""
PyTorch，把我的模型 net 变成一个多卡并行版本，让它在 devices 列表里的所有显卡上同时跑，并把主控制节点设在第 0 号显卡上，这样训练起来就快多了"
"""    


    # 使用 Adam 优化器 (决定了模型在每次发现错误后，如何去调整自己的脑部参数)
    trainer = torch.optim.Adam(net.parameters(), lr=0.01)
    
    step, timer = 0, d2l.Timer()
    # 动画画图工具：在训练时实时画出 mlm 和 nsp 的误差下降曲线
    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:
        # 从流水线 (DataLoader) 里源源不断地抓取批次数据 (每次 512 对句子)
        for tokens_X, segments_X, valid_lens_x, pred_positions_X,\
            mlm_weights_X, mlm_Y, nsp_y in train_iter:
            
            # 【关键操作】把数据从慢吞吞的内存(CPU)搬运到飞快的显卡(GPU)上
            tokens_X = tokens_X.to(devices[0])
            segments_X = segments_X.to(devices[0])
            valid_lens_x = valid_lens_x.to(devices[0])
            pred_positions_X = pred_positions_X.to(devices[0])
            mlm_weights_X = mlm_weights_X.to(devices[0])
            mlm_Y, nsp_y = mlm_Y.to(devices[0]), nsp_y.to(devices[0])
            
            # 1. 擦黑板：清空上一次计算残留的梯度 (如果不清空，误差会越积越多)
            trainer.zero_grad()
            timer.start()
            
            # 2. 判卷子：计算这 512 道题的总误差 l (包含 mlm_l 和 nsp_l)
            mlm_l, nsp_l, 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)
                
            # 3. 反向传播：把误差信号顺着网线传回去，告诉每一层神经元"你错哪了"
            l.backward()
            
            # 4. 纠正错误：优化器根据刚才传回来的错误信号，稍微修改一下模型的参数
            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 loss {metric[0] / metric[3]:.3f}, '
          f'NSP loss {metric[1] / metric[3]:.3f}')
    print(f'{metric[2] / timer.sum():.1f} sentence pairs/sec on '
          f'{str(devices)}')

train_bert(train_iter, net, loss, len(vocab), devices, 50)

50step，还没有收敛:

Batch (批次): 我们通常不会一次性把所有数据都塞给模型（内存会爆），而是把数据分成一小块一小块的，这一小块就是一个 Batch。比如 batch_size=32，意思是每次拿 32 句话给模型训练。
Step (步骤/迭代次): 模型每处理完一个 Batch，更新一次参数，就叫做一个 Step（或者一个 Iteration）。
Epoch (轮次): 当模型把整个训练集（所有的数据）都完整地看过一遍，才叫做一个 Epoch (一轮)。

一个 Epoch 包含很多个 Steps。

用Bert表示文本:

在预训练BERT之后，我们可以用它来表示单个文本、文本对或其中的任何词元。下面的函数返回tokens_a和tokens_b中所有词元的BERT（net）表示。

python 复制代码

import torch

# ==========================================
# 1. 定义特征提取辅助函数
# ==========================================
def get_bert_encoding(net, tokens_a, tokens_b=None):
    """
    这个函数的作用是：把人类的单词丢进训练好的 BERT，榨取它们的"高维向量表示"
    """
    # 1. 组装输入：自动帮我们加上 '<cls>' 和 '<sep>'，并生成对应的片段号码牌 segments
    tokens, segments = d2l.get_tokens_and_segments(tokens_a, tokens_b)
    
    # 2. 查字典并加上批次维度：
    # vocab[tokens] 把单词变数字。
    # .unsqueeze(0) 是在最前面强行加一个维度。因为 BERT 默认接收的是一个"批次(Batch)"，
    # 我们这里只有 1 道题，所以要把形状从 [长度] 变成 [1, 长度]
    token_ids = torch.tensor(vocab[tokens], device=devices[0]).unsqueeze(0)
    segments = torch.tensor(segments, device=devices[0]).unsqueeze(0)
    
    # 3. 记录有效长度 (因为这里没有用到 '<pad>' 填充，所以有效长度就是整个句子的长度)
    valid_len = torch.tensor(len(tokens), device=devices[0]).unsqueeze(0)
    
    # 4. 前向传播，提取特征！
    # 还记得 BERT 返回三个东西吗？(特征矩阵, MLM预测, NSP预测)
    # 我们这里只是想看看词的特征长什么样，不需要它做预测，所以用 `_` 把后两个扔掉
    encoded_X, _, _ = net(token_ids, segments, valid_len)
    
    # 返回提取到的特征矩阵，形状大概是 (1, 句子总长度, 128)
    return encoded_X  


# ==========================================
# 2. 测试场景一：单句输入 (crane 表示"鹤")
# ==========================================
# 句子意思："一只鹤正在飞"
tokens_a = ['a', 'crane', 'is', 'flying']
encoded_text = get_bert_encoding(net, tokens_a)

# 经过上面的函数，实际进入 BERT 的词元变成了：
# 索引:   0       1       2       3        4        5
# 词元: '<cls>', 'a', 'crane', 'is', 'flying', '<sep>'

# 提取 <cls> 的特征：它在第 0 个位置，代表了"整句话的宏观大意"
encoded_text_cls = encoded_text[:, 0, :]

# 提取 'crane' (鹤) 的特征：它在第 2 个位置
encoded_text_crane = encoded_text[:, 2, :]

# 打印看看它们的形状，以及 'crane' 这个词特征向量的前 3 个数字
print("单句整体形状:", encoded_text.shape) 
print("<cls>形状:", encoded_text_cls.shape)
print("'crane'(鹤)的前3个

测试场景

考虑"a crane is flying"这句话。回想一下 :numref:subsec_bert_input_rep中讨论的BERT的输入表示。插入特殊标记"<cls>"（用于分类）和"<sep>"（用于分隔）后，BERT输入序列的长度为6。因为零是"<cls>"词元，encoded_text[:, 0, :]是整个输入语句的BERT表示。为了评估一词多义词元"crane"，我们还打印出了该词元的BERT表示的前三个元素。

现在考虑一个句子"a crane driver came"和"he just left"。类似地，encoded_pair[:, 0, :]是来自预训练BERT的整个句子对的编码结果。注意，多义词元"crane"的前三个元素与上下文不同时的元素不同。这支持了BERT表示是上下文敏感的。

python 复制代码

# ==============================================================
# 场景一：单句输入 (在这个语境下，'crane' 表示 "鹤")
# ==============================================================
# 假设我们输入一句话："一只鹤正在飞" (a crane is flying)
tokens_a = ['a', 'crane', 'is', 'flying']

# 将分词后的列表送入 `get_bert_encoding` 函数。
# 这个函数内部会负责添加特殊词元、转换为 ID、并送入预训练的 BERT (net) 模型中。
encoded_text = get_bert_encoding(net, tokens_a)

# --------------------------------------------------------------
# 内部发生了什么？
# 经过处理后，实际喂给 BERT 模型的词元序列变成了这样：
# 索引：   0       1        2        3        4         5
# 词元： '<cls>', 'a',   'crane',  'is',  'flying', '<sep>'
# 
# 注意：
# '<cls>' (Classification) 是强制加在最前面的。
# '<sep>' (Separator) 是加在句末，表示这半句话结束了。
# 这就是为什么输入 4 个词，最终序列长度变成了 6。
# --------------------------------------------------------------

# 提取 <cls> 词元的特征表示。
# [:, 0, :] 的意思是：取所有批次(batch)、第 0 个位置、所有的隐藏层特征(如128维)。
# 为什么要取第 0 个位置？
# 因为 BERT 规定，<cls> 位置输出的向量，融合了整句话的全局信息。
# 如果你要给这句话做情感分类（积极/消极），通常就是把这个 `encoded_text_cls` 拿去接个分类器。
encoded_text_cls = encoded_text[:, 0, :]

# 提取 'crane' (鹤) 这个特定词元的特征表示。
# 对照上面的索引表，'crane' 排在第 2 个位置（因为前面插了个<cls>），所以取索引 2。
encoded_text_crane = encoded_text[:, 2, :]

# 打印出它们的形状，以及 'crane' 这个词特征向量的前 3 个数值，方便我们观察。
# encoded_text 的形状大概是 (batch_size, 6, hidden_size)
print("单句整体形状 (包含所有词元):", encoded_text.shape) 
print("<cls> 词元形状 (代表整句话):", encoded_text_cls.shape)
print("'crane'(鹤) 特征向量的前3个值:", encoded_text_crane[0][:3])

"""
(torch.Size([1, 6, 128]),
 torch.Size([1, 128]),
 tensor([-0.5007, -1.0034,  0.8718], device='cuda:0', grad_fn=<SliceBackward0>))

[1, 6, 128] 告诉你：在这个 1 句话里，总共有 6 个词元，BERT 为这 6 个词元中的每一个，都计算出了一个长度为 128 的特征向量。
"""

# ==============================================================
# 场景二：句子对输入 (在这个语境下，'crane' 表示 "起重机")
# ==============================================================
# 假设我们输入两句话，比如在做问答系统或者自然语言推理。
# 第一句："一个起重机司机来了" (a crane driver came)
# 第二句："他刚离开" (he just left)
tokens_a, tokens_b = ['a', 'crane', 'driver', 'came'], ['he', 'just', 'left']

# 同样送入函数进行编码。这次有 tokens_a 和 tokens_b 两个输入。
encoded_pair = get_bert_encoding(net, tokens_a, tokens_b)

# --------------------------------------------------------------
# 内部发生了什么？
# 对于句子对，BERT 会把它们拼起来，并用 <sep> 隔开。词元序列变成：
# 索引：   0       1       2         3        4         5       6      7       8       9
# 词元： '<cls>', 'a',  'crane', 'driver', 'came', '<sep>', 'he', 'just', 'left', '<sep>'
#
# 注意：两个句子被拼成了长度为 10 的单个序列。
# --------------------------------------------------------------

# 再次提取 <cls> 词元的特征表示。
# 虽然这次输入的是两句话，但第 0 个位置的 <cls> 依然尽职尽责。
# 此时，这个 `encoded_pair_cls` 向量融合了**这两句话拼接在一起后的整体全局信息**。
# 它可以用来判断这两句话是什么关系（比如：是不是因果关系？是不是前后矛盾？）。
encoded_pair_cls = encoded_pair[:, 0, :]

# 再次提取 'crane' (此时是"起重机"的意思) 这个词元的特征表示。
# 在新的拼接序列中，'crane' 依然恰巧在第 2 个位置，所以取索引 2。
encoded_pair_crane = encoded_pair[:, 2, :]

# 打印观察结果。
print("\n--- 切换到句子对 ---")
print("句子对整体形状 (长度变长了):", encoded_pair.shape)
print("句子对 <cls> 形状:", encoded_pair_cls.shape)

# 【核心重点】：
# 当你对比这里打印出的 `encoded_pair_crane[0][:3]` 和上面场景一的 `encoded_text_crane[0][:3]` 时，
# 你会发现**这前三个数字（以及后面的所有数字）是不一样的！**
print("'crane'(起重机) 特征向量的前3个值 (注意与上面不同):", encoded_pair_crane[0][:3])

"""
(torch.Size([1, 10, 128]),
 torch.Size([1, 128]),
 tensor([ 0.5101, -0.4041, -1.2749], device='cuda:0', grad_fn=<SliceBackward0>))
"""
# 为什么不一样？
# 在以前简单的模型里，单词 "crane" 的数字表示是固定的死板的。
# 而 BERT 是"上下文敏感"的 (context-sensitive)。
# 在第一句话里，BERT 看到旁边有 "flying"(飞)，就调整了 "crane" 的向量，让它偏向"鸟类"的特征。
# 在第二句话里，BERT 看到旁边有 "driver"(司机)，就调整了 "crane" 的向量，让它偏向"机器"的特征。
# 这样，BERT 就巧妙地解决了一词多义的问题。

拓展:

有打印结果，但是结果不是人可阅读或者理解的

其次如果是中文:

中文用句号或者nltk拆分句子，然后用jieba或其他的做分词当作token，接下来就一样训练了