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

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"]  # 返回文本和标签 

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 

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} 保存成功")