BERT训练环节（代码实现）

1.代码实现

#导包
import torch
from torch import nn
import dltools

#加载数据需要用到的声明变量
batch_size, max_len = 1, 64
#获取训练数据迭代器、词汇表
train_iter, vocab = dltools.load_data_wiki(batch_size, max_len)

#其余都是二维数组
#tokens, segments, valid_lens（一维）, pred_position, mlm_weights, mlm, nsp（一维）对应每条数据i中包含的数据
for i in train_iter:  #遍历迭代器
    break   #只遍历一条数据

[tensor([[    3,    25,     0,  4993,     0,    24,     4,    26,    13,     2,
            158,    20,     5,    73,  1399,     2,     9,   813,     9,   987,
             45,    26,    52,    46,    53,   158,     2,     5,  3140,  5880,
              9,   543,     6,  6974,     2,     2,   315,     6,     8,     5,
           8698,     8, 17229,     9,   308,     2,     4,     1,     1,     1,
              1,     1,     1,     1,     1,     1,     1,     1,     1,     1,
              1,     1,     1,     1]]),
 tensor([[0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0,
          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]),
 tensor([47.]),
 tensor([[ 9, 15, 26, 32, 34, 35, 45,  0,  0,  0]]),
 tensor([[1., 1., 1., 1., 1., 1., 1., 0., 0., 0.]]),
 tensor([[ 484, 1288,   20,    6, 2808,    9,   18,    0,    0,    0]]),
 tensor([0])]

#创建BERT网络模型
net = dltools.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)
#调用设备上的GPU
devices = dltools.try_all_gpus()
#损失函数对象
loss = nn.CrossEntropyLoss()   #多分类问题，使用交叉熵

#@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):
    #前向传播
    #获取遮蔽词元的预测结果、下一个句子的预测结果
    _, mlm_Y_hat, nsp_Y_hat = net(tokens_X, segments_X, valid_lens_x.reshape(-1), pred_positions_X)
    #计算遮蔽语言模型的损失
    mlm_l = loss(mlm_Y_hat.reshape(-1, vocab_size), mlm_Y.reshape(-1)) * mlm_weights_X.reshape(-1,1)
    mlm_l = mlm_l.sum() / (mlm_weights_X.sum() + 1e-8)   #MLM损失函数的归一化版本   #加一个很小的数1e-8，防止分母为0，抵消上一行代码乘以的数值
    #计算下一个句子预测任务的损失
    nsp_l = loss(nsp_Y_hat, nsp_y)
    l = mlm_l + nsp_l
    return mlm_l, nsp_l, l  

def train_bert(train_iter, net, loss, vocab_size, devices, num_steps):  #文本词元样本量太多，全跑完花费的时间太多，若num_steps=1在BERT中表示，跑了1个batch_size
    net = nn.DataParallel(net, device_ids=devices).to(devices[0])  #调用设备的GPU
    trainer = torch.optim.Adam(net.parameters(), lr=0.01)   #梯度下降的优化算法Adam
    step, timer = 0, dltools.Timer()  #设置计时器
    #调用画图工具
    animator = dltools.Animator(xlabel='step', ylabel='loss', xlim=[1, num_steps], legend=['mlm', 'nsp'])
    #遮蔽语言模型损失的和， 下一句预测任务损失的和， 句子对的数量， 计数
    metric = dltools.Accumulator(4)  #Accumulator类被设计用来收集和累加各种指标（metric）
    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:
            #将遍历的数据发送到设备上
            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])
            
            #梯度清零
            trainer.zero_grad()
            timer.start()  #开始计时
            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)
            l.backward()  #反向传播
            trainer.step()  #梯度更新
            metric.add(mlm_l, nsp_l, tokens_X.shape[0], l)  #累积的参数指标
            timer.stop() #计时停止
            animator.add(step + 1, (metric[0] / metric[3], metric[1] / metric[3]))  #画图的
            step += 1  #训练完一个batch_size,就+1
            if step == num_steps:  #若步数与预设的训练步数相等
                num_steps_reached = True   #判断标志改为True
                break  #退出while循环
    
    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, 500)

def get_bert_encoding(net, tokens_a, tokens_b=None):
    tokens, segments = dltools.get_tokens_and_segments(tokens_a, tokens_b)
    token_ids = torch.tensor(vocab[tokens], device=devices[0]).unsqueeze(0)  #unsqueeze(0)增加一个维度
    segments = torch.tensor(segments, device=devices[0]).unsqueeze(0)  
    valid_len = torch.tensor(len(tokens), device=devices[0]).unsqueeze(0)
    endoced_X, _, _ = net(token_ids, segments, valid_len)
    return endoced_X

tokens_a = ['a', 'crane', 'is', 'flying']
encoded_text = get_bert_encoding(net, tokens_a)
# 词元：'<cls>','a','crane','is','flying','<sep>'
encoded_text_cls = encoded_text[:, 0, :]
encoded_text_crane = encoded_text[:, 2, :]
encoded_text.shape, encoded_text_cls.shape, encoded_text_crane[0][:3]

(torch.Size([1, 6, 128]),
 torch.Size([1, 128]),
 tensor([-0.5872, -0.0510, -0.7376], device='cuda:0', grad_fn=<SliceBackward0>))

encoded_text_crane

 

tensor([[-5.8725e-01, -5.0994e-02, -7.3764e-01, -4.3832e-02,  9.2467e-02,
          1.2745e+00,  2.7062e-01,  6.0271e-01, -5.5055e-02,  7.5122e-02,
          4.4872e-01,  7.5821e-01, -6.1558e-02, -1.2549e+00,  2.4479e-01,
          1.3132e+00, -1.0382e+00, -4.7851e-03, -6.3590e-01, -1.3180e+00,
          5.2245e-02,  5.0982e-01,  7.4168e-02, -2.2352e+00,  7.4425e-02,
          5.0371e-01,  7.2120e-02, -4.6384e-01, -1.6588e+00,  6.3987e-01,
         -6.4567e-01,  1.7187e+00, -6.9696e-01,  5.6788e-01,  3.2628e-01,
         -1.0486e+00, -7.2610e-01,  5.7909e-02, -1.6380e-01, -1.2834e+00,
          1.6431e+00, -1.5972e+00, -4.5678e-03,  8.8022e-02,  5.5931e-02,
         -7.2332e-02, -4.9313e-01, -4.2971e+00,  6.9757e-01,  7.0690e-02,
         -1.8613e+00,  2.0366e-01,  8.9868e-01, -3.4565e-01,  9.6776e-02,
          1.3699e-02,  7.1410e-01,  5.4820e-01,  9.7358e-01, -8.1038e-01,
          2.6216e-01, -5.7850e-01, -1.1969e-01, -2.5277e-01, -2.0046e-01,
         -1.6718e-01,  5.5540e-01, -1.8172e-01, -2.5639e-02, -6.0961e-01,
         -1.1521e-03, -9.2973e-02,  9.5226e-01, -2.4453e-01,  9.7340e-01,
         -1.7908e+00, -2.9840e-02,  2.3087e+00,  2.4889e-01, -7.2734e-01,
          2.1827e+00, -1.1172e+00, -7.0915e-02,  2.5138e+00, -1.0356e+00,
         -3.7332e-02, -5.6668e-01,  5.2251e-01, -5.0058e-01,  1.7354e+00,
          4.0760e-01, -1.2982e-01, -7.0230e-01,  3.1563e+00,  1.8754e-01,
          2.0220e-01,  1.4500e-01,  2.3296e+00,  4.5522e-02,  1.1762e-01,
          1.0662e+00, -4.0858e+00,  1.6024e-01,  1.7885e+00, -2.7034e-01,
         -1.6869e-01, -8.7018e-02, -4.2451e-01,  1.1446e-01, -1.5761e+00,
          7.6947e-02,  2.4336e+00,  4.5346e-02, -6.5078e-02,  1.4203e+00,
          3.7165e-01, -7.9571e-01, -1.3515e+00,  4.1511e-02,  1.3561e-01,
         -3.3006e+00,  1.4821e-01,  1.3024e-01,  1.9966e-01, -8.5910e-01,
          1.4505e+00,  7.6774e-02,  9.3771e-01]], device='cuda:0',
       grad_fn=<SliceBackward0>)

tokens_a, tokens_b = ['a', 'crane', 'driver', 'came'], ['he', 'just', 'left']
encoded_pair = get_bert_encoding(net, tokens_a, tokens_b)
# 词元：'<cls>','a','crane','driver','came','<sep>','he','just', 'left','<sep>'
encoded_pair_cls = encoded_pair[:, 0, :]
encoded_pair_crane = encoded_pair[:, 2, :]
encoded_pair.shape, encoded_pair_cls.shape, encoded_pair_crane[0][:3]

 

(torch.Size([1, 10, 128]),
 torch.Size([1, 128]),
 tensor([-0.4637, -0.0569, -0.6119], device='cuda:0', grad_fn=<SliceBackward0>))