15 大模型 LLM 微调过程中保存与加载模型

在大模型的微调过程中，保存 和 加载 模型是非常重要的一环。它不仅可以帮助我们保存训练好的模型以供后续使用，还可以在需要时快速恢复训练状态，以便进行进一步的微调或评估。在本篇文章中，我们将深入探讨如何有效地保存和加载大模型。

保存模型

在微调过程中，保存模型的主要目的是为了防止数据丢失，以及在达到一定的训练轮次后，可以对模型状态进行持久化。

使用 Hugging Face Transformers 库保存模型

假设我们在进行微调的过程使用的是 Hugging Face 的 Transformers 库，可以通过如下代码保存模型：

from transformers import Trainer, TrainingArguments

# 假设 trainer 是已经初始化好的 Trainer 对象
trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="./results",
        save_total_limit=2,  # 仅保留最近的两个模型
        save_steps=500,      # 每500步保存一次模型
    ),
    train_dataset=train_dataset,
)

# 开始训练并自动保存模型
trainer.train()

在上述代码中，output_dir 用于指定模型保存的路径，save_steps 表示每经过 500 步训练就保存一次模型。使用 save_total_limit 参数可以控制保存的模型数量，避免占用过多的存储空间。

保存模型权重

除去完整的模型和配置之外，有时我们可能只需保存模型的权重，可以使用以下代码：

model.save_pretrained("./model_weights")
tokenizer.save_pretrained("./model_weights")

这将会把模型的权重和 tokenizer 保存到指定的文件夹中。

加载模型

在微调的后续过程中，或者在需要进行模型评估和推理时，加载之前保存的模型是必不可少的。我们可以通过类似的 Transformers 方法来实现模型的加载。

加载完整模型和Tokenizer

可以 воспользоваться следующим кодом для загрузки модели и токенизатора из ранее сохраненной папки：

from transformers import AutoModelForSequenceClassification, AutoTokenizer

# 加载模型
model = AutoModelForSequenceClassification.from_pretrained("./model_weights")
tokenizer = AutoTokenizer.from_pretrained("./model_weights")

加载特定权重

如果你只想加载模型的权重，可以单独执行以下命令：

model.load_state_dict(torch.load("./model_weights/pytorch_model.bin"))

案例演示

让我们举一个简单的例子，展示如何在微调完成后进行模型的保存与加载。

微调过程

假设我们微调了一个BERT模型进行情感分析，完整代码如下：

from transformers import BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments

# 1. 加载模型和tokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForSequenceClassification.from_pretrained('bert-base-uncased')

# 2. 准备数据集（省略具体数据预处理步骤）

# 3. 设置训练参数
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=8,
    save_steps=500,
)

# 4. 创建 Trainer 对象并开始训练
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
)

# 5. 开始训练并自动保存
trainer.train()

加载模型进行评估

训练完成后，我们想要进行模型的评估，可以使用如下代码加载模型并进行预测：

# 加载保存的模型和tokenizer
model = BertForSequenceClassification.from_pretrained('./results')
tokenizer = BertTokenizer.from_pretrained('./results')

# 进行推理（假设有一段新的文本需要预测）
inputs = tokenizer("I love using transformers!", return_tensors="pt")
outputs = model(**inputs)

# 解析模型输出
logits = outputs.logits
predicted_class = logits.argmax(dim=-1)
print(f"预测的情感类别: {predicted_class.item()}")

在这个例子中，我们完成了模型的微调并保存了其状态，然后成功加载该模型并进行了评估。通过以上方法，我们可以灵活地保存和加载模型，以便在实际应用中最大化模型的利用率。

总结

本节的重点是如何在大模型 LLM 微调过程中有效地 保存 和 加载 模型。合理的保存和加载策略不仅有助于节省训练时间，还有助于更好地利用已有的计算资源。在下一篇文章中，我们将探讨如何设置评估指标，以便对微调后的模型进行有效的评估和性能测试。