【动手学习pytorch笔记】37.3 BERT预训练

最新推荐文章于 2025-08-18 15:35:34 发布

原创

最新推荐文章于 2025-08-18 15:35:34 发布 · 806 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #学习 #bert #nlp #深度学习

该代码示例展示了如何在PyTorch中实现BERT预训练模型的训练过程，包括遮蔽语言模型（MLM）和下一句预测（NSP）任务。由于模型的复杂性，通常需要大量的步骤（如100,000以上）来获得良好的预训练效果。训练完成后，可以使用模型对任意句子进行特征抽取。

BERT预训练

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

以下都是很常规的

batch_size, max_len = 512, 64
train_iter, vocab = d2l.load_data_wiki(batch_size, max_len)

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()
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(-