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参数，包括原理说明和完整代码实现。这种方法可以扩展到其他参数的修改，为定制化预训练模型提供了参考。关键点在于正确修改配置、替换相应层并合理初始化参数。