pytorch入门 - 修改huggingface大模型配置参数

介绍

Hugging Face的Transformers库提供了大量预训练模型，但有时我们需要修改这些模型的默认参数来适应特定任务。

本文将详细介绍如何修改BERT模型的最大序列长度(max_position_embeddings)参数，并解释相关原理和实现细节。

原理

BERT等Transformer模型对输入序列长度有固定限制，这主要由位置编码(position embeddings)决定。

原始BERT-base-chinese模型的max_position_embeddings为512，意味着它最多只能处理512个token的输入。当我们需要处理更长的文本时，必须修改这一参数。

修改过程涉及三个关键步骤：

1. 调整模型配置中的max_position_embeddings值
2. 替换位置嵌入层(position_embeddings)为新尺寸
3. 初始化新位置嵌入层的权重(复制原有权重，其余随机初始化)

实现代码详解

下面我们逐行分析实现代码：

1. 数据集准备 (news_finetuing_data_set.py)

from datasets import load_dataset, load_from_disk
from torch.utils.data import Dataset
class MyDataset(Dataset):
 def __init__(self, split):
 # 指定CSV文件路径，支持train/test/validation三种分割
 data_file = rf"cache\datasets\csv\THUCNewsText\{split}.csv"
 self.dataset = load_dataset(
 "csv",
 data_files={split: data_file},
 split=split if split in ["train", "test", "validation"] else "train",
 )
 def __len__(self):
 return len(self.dataset) # 返回数据集样本数量
 def __getitem__(self, idx):
 return self.dataset[idx]["text"], self.dataset[idx]["label"] # 返回文本和标签

2. 模型修改 (news_finetuing_net.py)

from transformers import BertModel, BertConfig
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 1. 加载预训练模型和配置
model = BertModel.from_pretrained(
 "bert-base-chinese", cache_dir="./cache/bertbasechinese"
).to(device)
# 2. 修改max_position_embeddings配置
model.config.max_position_embeddings = 1500
# 3. 替换position_embeddings层
old_embeddings = model.embeddings.position_embeddings
new_embeddings = torch.nn.Embedding(1500, old_embeddings.embedding_dim)
# 拷贝原有权重
num = min(old_embeddings.weight.size(0), 1500)
new_embeddings.weight.data[:num, :] = old_embeddings.weight.data[:num, :]
model.embeddings.position_embeddings = new_embeddings
# 4. 冻结除position_embeddings外的所有参数
for name, param in pretrained.named_parameters():
 if "embeddings.position_embeddings" in name:
 param.requires_grad = True
 else:
 param.requires_grad = False
class Model(torch.nn.Module):
 def __init__(self):
 super(Model, self).__init__()
 self.classifier = torch.nn.Linear(768, 10) # 添加分类头
 def forward(self, input_ids, attention_mask, token_type_ids):
 position_ids = (
 torch.arange(input_ids.size(1), dtype=torch.long, device=input_ids.device)
 .unsqueeze(0)
 .expand_as(input_ids)
 )
 outputs = pretrained(
 input_ids=input_ids,
 attention_mask=attention_mask,
 token_type_ids=token_type_ids,
 position_ids=position_ids,
 )
 cls_output = outputs.last_hidden_state[:, 0] # 取[CLS] token的输出
 out = self.classifier(cls_output)
 return out

3. 训练过程 (news_finetuing_train.py)

import torch
from news_finetuing_data_set import MyDataset
from torch.utils.data import DataLoader
from news_finetuing_net import Model
from transformers import BertTokenizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
EPOCH = 100
# 加载分词器并设置最大长度
token = BertTokenizer.from_pretrained(
 "bert-base-chinese",
 cache_dir="./cache/tokenizer/bert-base-chinese",
)
token.model_max_length = 1500 # 设置分词器最大长度
def collate_fn(batch):
 # 数据处理函数
 sentes = [item[0] for item in batch]
 labels = [item[1] for item in batch]
 data = token.batch_encode_plus(
 sentes,
 truncation=True,
 padding="max_length",
 max_length=1500,
 return_tensors="pt",
 return_length=True,
 )
 # 返回模型需要的各种输入
 return (
 data["input_ids"],
 data["attention_mask"],
 data["token_type_ids"],
 torch.LongTensor(labels),
 )
# 创建数据集和DataLoader
train_dataset = MyDataset(split="train")
val_dateset = MyDataset(split="validation")
train_loader = DataLoader(
 train_dataset,
 batch_size=32,
 shuffle=True,
 drop_last=True,
 collate_fn=collate_fn,
)
val_loader = DataLoader(
 val_dateset,
 batch_size=32,
 shuffle=False,
 drop_last=True,
 collate_fn=collate_fn,
)
# 训练主循环
if __name__ == "__main__":
 model = Model().to(device)
 optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
 loss_func = torch.nn.CrossEntropyLoss()
 for epoch in range(EPOCH):
 model.train()
 for step, (input_ids, attention_mask, token_type_ids, labels) in enumerate(train_loader):
 # 数据移动到设备
 input_ids = input_ids.to(device)
 attention_mask = attention_mask.to(device)
 token_type_ids = token_type_ids.to(device)
 labels = labels.to(device)
 # 前向传播和反向传播
 outputs = model(input_ids, attention_mask, token_type_ids)
 loss = loss_func(outputs, labels)
 optimizer.zero_grad()
 loss.backward()
 optimizer.step()
 # 打印训练信息
 if step % 5 == 0:
 out = outputs.argmax(dim=1)
 acc = (out == labels).sum().item() / len(labels)
 print(f"Epoch: {epoch + 1}/{EPOCH}, Step: {step + 1}/{len(train_loader)}, Loss: {loss.item():.4f}, Acc: {acc:.4f}")
 # 保存模型
 torch.save(model.state_dict(), f"./model/news_finetuning_epoch_{epoch}.pth")
 print(f"epoch {epoch} 保存成功")

关键点解释

模型修改部分

1. model.config.max_position_embeddings = 1500 - 修改配置中的最大位置嵌入数
2. 创建新的位置嵌入层时，我们保留了原始嵌入维度(embedding_dim)，只扩展了位置数量
3. 权重初始化策略是复制原有512个位置的权重，剩余位置使用随机初始化

训练策略

1. 我们冻结了除位置嵌入外的所有BERT参数，只训练位置嵌入和新添加的分类头
2. 这种策略在长文本微调中很常见，可以防止过拟合

数据处理

1. 分词器也需要设置model_max_length以匹配新的序列长度
2. collate_fn函数确保所有输入都被填充/截断到1500的长度

总结

本文详细介绍了如何修改Hugging Face模型的max_position_embeddings参数，包括原理说明和完整代码实现。这种方法可以扩展到其他参数的修改，为定制化预训练模型提供了参考。关键点在于正确修改配置、替换相应层并合理初始化参数。