1.代码实现
python">#导包
import torch
from torch import nn
import dltoolspython">#加载数据需要用到的声明变量
batch_size, max_len = 1, 64
#获取训练数据迭代器、词汇表
train_iter, vocab = dltools.load_data_wiki(batch_size, max_len)python">#其余都是二维数组
#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])]
python">#创建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() #多分类问题，使用交叉熵python">#@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_lreturn mlm_l, nsp_l, l python">def train_bert(train_iter, net, loss, vocab_size, devices, num_steps): #文本词元样本量太多，全跑完花费的时间太多，若num_steps=1在BERT中表示，跑了1个batch_sizenet = nn.DataParallel(net, device_ids=devices).to(devices[0]) #调用设备的GPUtrainer = torch.optim.Adam(net.parameters(), lr=0.01) #梯度下降的优化算法Adamstep, 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,就+1if step == num_steps: #若步数与预设的训练步数相等num_steps_reached = True #判断标志改为Truebreak #退出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)}')python">train_bert(train_iter, net, loss, len(vocab), devices, 500) 
python">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_Xpython">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>))
python">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>)
python">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>))
