文章目录
概要
摘本文详细介绍了基于BERT的中文掩码语言模型(MLM)实现过程。主要包括：1) 使用Hugging Face的BertTokenizer和BertModel加载预训练模型；2) 自定义数据预处理函数collate_fn2，构建包含[MASK]的输入序列和对应标签；3) 设计包含线性输出层的MyModel模型架构；4) 采用AdamW优化器和交叉熵损失进行模型训练；5) 实现模型评估流程。
关键技术点包括：固定位置掩码处理、BERT参数冻结、CUDA加速等。实验结果表明，该方法能有效完成中文文本的掩码词预测任务。代码完整展示了从数据加载到模型训练的完整流程。
整体架构流程
自定义函数
获取数据集
定义模型
模型训练
模型评估
技术细节
1、导包
import torch
import torch.nn as nn
from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers import BertTokenizer, BertModel
from transformers import AdamW
import torch.optim as optim
import time
from tqdm import tqdm
2、使用CUDA、加载分词器和模型
device = 'cuda' if torch.cuda.is_available() else "cpu"
# 分词器
bert_tokenizer = BertTokenizer.from_pretrained('../dm04fasttext/models/bert-base-chinese')
# 模型
bert_model = BertModel.from_pretrained('../dm04fasttext/models/bert-base-chinese')
bert_model = bert_model.to(device)
print(bert_model)3、自定义函数collate_fn2
对一个批次的文本数据进行预处理，构造BERT模型的输入和标签。具体包括：
-
编码文本：使用BERT的分词器将文本转为ID序列，并统一填充/截断到固定长度（32）。
-
构造标签：提取每个样本第16个位置的词作为标签（模拟BERT的掩码语言任务，MLM）。
-
掩码处理：将原始文本中第16个词替换为
[MASK]，让模型预测该位置的原始词（标签）。
最终返回格式化后的输入（input_ids, token_type_ids, attention_mask）和对应的标签（labels
关键字含义
input_ids：编码后的 token IDs（其中第16个位置被[MASK]替换）。
token_type_ids：句子分段标记（通常用于区分句子）。
attention_mask：标记哪些位置是有效 token（非 padding）。
labels：被掩盖的第16个位置的原始 token ID（用于计算MLM损失）。
batch_encode_plus：对批次文本统一编码，生成以下字段：
input_ids：token 对应的词汇表 ID。
token_type_ids：句子分段标记（单句任务时通常全为0）。
clone()：创建副本，避免后续修改影响原始数据attention_mask：标记有效 token（1=真实 token，0=padding）
代码
def collate_fn2(data):
# print(f'data-->{data}')
# 获取每个样本的评论数据
sents = [value["text"] for value in data]
# 对上述的一个批次的样本数据进行编码
inputs = bert_tokenizer.batch_encode_plus(sents, padding='max_length',
truncation=True, max_length=32, return_tensors='pt')
# print(f'inputs-->{inputs}')
input_ids = inputs["input_ids"]
token_type_ids = inputs["token_type_ids"]
attention_mask = inputs["attention_mask"]
# input_ids-->[4, 32]
# 取出第16个位置的元素
labels = input_ids[:, 16].reshape(-1).clone()
# 将原始的每个样本的第16个词，进行mask掩盖
input_ids[:, 16] = bert_tokenizer.get_vocab()[bert_tokenizer.mask_token]
# input_ids[:, 16] = bert_tokenizer.mask_token_id
labels = torch.tensor(labels, dtype=torch.long)
return input_ids, token_type_ids, attention_mask, labels
返回格式：
input_ids：掩盖后的 token IDs（形状[batch_size, 32]）。
token_type_ids：句子分段标记（形状[batch_size, 32]）。
attention_mask：注意力掩码（形状[batch_size, 32]）。
labels：被掩盖的原始 token ID（形状[batch_size]）
4、获取数据集
关键字含义：
drop_last=True：丢弃最后不足一个批次的数据（避免批次大小不一致）
代码：
def get_dataloader():
# 基于load_dataset方法读取原始的数据
train_dataset = load_dataset('csv', data_files='../day09/data/train.csv', split='train')
# print(f'train_dataset--》{train_dataset}')
# print(f'train_dataset--》{train_dataset[0]}')
# 只保留原始评论文本大于32的样本
new_train_dataset = train_dataset.filter(lambda x: len(x["text"]) > 32)
# print(f'new_train_dataset-->{new_train_dataset}')
# 实例化dataloader
train_dataloader = DataLoader(dataset=new_train_dataset,
batch_size=4,
collate_fn=collate_fn2,
shuffle=True,
drop_last=True)
return train_dataloader
输出结果
print("train_dataset",train_dataset) 输出如下train_dataset Dataset({
features: ['label', 'text'],
num_rows: 9600
})
print("train_dataset",train_dataset[0]) 输出如下train_dataset {'label': 1, 'text': '选择珠江花园的原因就是方便，有电动扶梯直接到达海边，周围餐馆、食廊、商场、超市、摊位一应俱全。酒店装修一般，但还算整洁。 泳池在大堂的屋顶，因此很小，不过女儿倒是喜欢。 包的早餐是西式的，还算丰富。 服务吗，一般'}
print("new_train_dataset",new_train_dataset) 输出如下new_train_dataset Dataset({
features: ['label', 'text'],
num_rows: 9035
})
5、定义模型
关键字含义：
关键字/代码 作用 nn.ModulePyTorch模型的基类，所有自定义模型需继承此类。 nn.Linear全连接层，将输入特征映射到输出空间（此处为词汇表大小）。 with torch.no_grad():禁用梯度计算，用于冻结预训练模型（如BERT）的参数。 bert_model预训练的BERT模型，用于提取文本特征。 last_hidden_stateBERT最后一层的隐藏状态（形状 [batch_size, seq_len, hidden_size]）。pooler_output[CLS]标记的聚合表示，通常用于句子分类任务（此处未使用）。 [:, 16]切片操作，选择所有样本的第16个token的特征。
代码：
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.out = nn.Linear(768, bert_tokenizer.vocab_size)
def forward(self, input_ids, token_type_ids, attention_mask):
with torch.no_grad():
bert_output = bert_model(input_ids=input_ids,
token_type_ids=token_type_ids,
attention_mask=attention_mask)
# print(f'bert_output--》{bert_output}')
# # last_hidden_state:[4, 32, 768]; pooler_output:[4, 768]
# print(f'bert_output--》{bert_output.last_hidden_state.shape}')
# print(f'bert_output--》{bert_output.pooler_output.shape}')
# 将上述的bert的last_hidden_state结果送入输出层,output-->[4, 21128]
output = self.out(bert_output.last_hidden_state[:, 16])
return output
nn.Linear(768, bert_tokenizer.vocab_size)：
输入维度：
768（BERT隐藏层大小）。输出维度：
bert_tokenizer.vocab_size（词汇表大小，如21128）。作用：将BERT的特征映射到词汇表空间，用于预测被掩盖的token
6、模型训练
代码
def model2train():
# 1.加载dataset对象
train_dataset = load_dataset('csv', data_files='../dm04fasttext/data/train.csv', split='train')
new_train_dataset = train_dataset.filter(lambda x: len(x["text"]) > 32)
# 2. 实例化模型
my_model = MyModel()
my_model = my_model.to(device)
# 3. 实例化损失函数对象
entropy = nn.CrossEntropyLoss()
# 4. 实例化优化器对象
adamw = AdamW(my_model.parameters(), lr=5e-4)
# adamw = optim.AdamW(my_model.parameters(), lr=5e-4)
# 5. 进一步操作，让预训练模型bert的参数不更新
for parameter in bert_model.parameters():
parameter.requires_grad_(False)
# 6. 指定模型为训练模式（一般只要用到预训练模型）
# 查询模型（预训练模型bert）的参数量
total_params = sum(p.numel() for p in bert_model.parameters())
print(f'total_params-->{total_params}')
my_model.train()
# 7.定义训练的轮次
epochs = 2
# 8. 开始外部循环
for epoch in range(epochs):
# 定义开始时间
start_time = int(time.time())
# 实例化dataloader
train_dataloader = DataLoader(dataset=new_train_dataset,
batch_size=8,
collate_fn=collate_fn2,
shuffle=True,
drop_last=True)
# 9.内部迭代
for idx, (input_ids, token_type_ids, attention_mask, labels) in enumerate(tqdm(train_dataloader), start=1):
# print(f'input_ids--》{input_ids.shape}')
input_ids = input_ids.to(device)
token_type_ids = token_type_ids.to(device)
attention_mask = attention_mask.to(device)
labels = labels.to(device)
output = my_model(input_ids, token_type_ids, attention_mask)
# print(f'output-->{output}')
# print(f'labels--》{labels}')
# print(f'labels--》{len(labels)}')
# 计算损失
loss = entropy(output, labels)
# 梯度清零
adamw.zero_grad()
# 反向传播
loss.backward()
# 梯度更新
adamw.step()
# print(f'loss-->{loss}')
# 每间隔5步，打印训练日志
if idx % 5 == 0:
predict = torch.argmax(output, dim=1)
avg_acc = (predict == labels).sum().item() / len(labels)
print('轮次:%d 迭代数:%d 损失:%.6f 准确率%.3f 时间%d' \
% (epoch, idx, loss.item(), avg_acc, (int)(time.time()) - start_time))
# 每一轮都保存模型
torch.save(my_model.state_dict(), 'fill_mask_model_%d.bin' % (epoch+1))
7、模型评估
def model2test():
# 1. 加载dataset对象
test_dataset = load_dataset('csv', data_files='../dm04fasttext/data/test.csv', split='train')
new_test_dataset = test_dataset.filter(lambda x: len(x["text"]) > 32)
# 2. 实例化dataloader
test_dataloader = DataLoader(dataset=new_test_dataset,
batch_size=8,
collate_fn=collate_fn2,
shuffle=True,
drop_last=True)
# 2. 实例化模型
my_model = MyModel()
my_model.load_state_dict(torch.load('fill_mask_model_3.bin'))
my_model = my_model.to(device)
# 3. 设置为评估模式
my_model.eval()
# 4. 准备一些打印日志的参数
correct_num = 0 # 预测正确的样本个数
total_num = 0 # 预测样本的总个数
# 5. 开始预测
# 9.内部迭代
with torch.no_grad():
for idx, (input_ids, token_type_ids, attention_mask, labels) in enumerate(tqdm(test_dataloader), start=1):
# print(f'input_ids--》{input_ids.shape}')
input_ids = input_ids.to(device)
token_type_ids = token_type_ids.to(device)
attention_mask = attention_mask.to(device)
labels = labels.to(device)
output = my_model(input_ids, token_type_ids, attention_mask)
# print(f'output-->{output.shape}')
# 计算预测正确的个数
predict_tag = torch.argmax(output, dim=1)
# print(f'predict_tag-->{predict_tag}')
correct_num += (predict_tag == labels).sum().item()
total_num += len(labels)
# print(f'labels-->{labels}')
# 每隔5步打印一下测试的日志
if idx % 5 == 0:
print(f'平均准确率：{correct_num/total_num}', end=' ')
# 每个批次获取其中一个样本查验
sample = bert_tokenizer.decode(input_ids[0])
print(f'{sample}', end=' ')
sample_predict_tag = bert_tokenizer.decode(predict_tag[0])
sample_target_tag = bert_tokenizer.decode(labels[0])
print(sample_predict_tag, sample_target_tag)8、函数入口
if __name__ == '__main__':
model2train()
model2test()
扫一扫
864
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
