人工智能

微调大模型的示范代码

2025年2月27日 作者 unix2go

以下是一个简单的 Hugging Face Transformers 框架中，用于微调大语言模型（如 GPT-2）的示范代码。该代码包括数据准备、模型加载、训练以及保存模型的完整流程。你可以根据自己的需求（如漏洞挖掘、渗透分析）调整数据和模型。

微调大模型示范代码

以下代码以 GPT-2 为例，对一个自定义文本数据集进行微调。代码假设你已经安装了 transformers 和 datasets 库，以及一个支持 GPU 的环境（如 PyTorch）。

1. 安装所需库

pip install transformers datasets torch

2. 微调代码示例

import os
import torch
from datasets import load_dataset
from transformers import GPT2Tokenizer, GPT2LMHeadModel, Trainer, TrainingArguments

# 1. 加载数据集
def load_custom_dataset(dataset_path):
    """
    加载自定义数据集
    数据格式假设为一个纯文本文件，每行是一条样本。
    """
    dataset = load_dataset("text", data_files={"train": dataset_path})
    return dataset

# 替换为你的数据集路径
dataset_path = "data/train.txt"
dataset = load_custom_dataset(dataset_path)

# 2. 加载预训练模型和分词器
model_name = "gpt2"  # 可以替换为 "gpt-neo", "LLaMA", 或其他支持的模型
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)

# GPT-2 默认不包含 `pad_token`，需要手动添加
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# 3. 数据预处理
def preprocess_function(examples):
    """
    将文本数据转化为模型所需的输入格式（tokenized）。
    """
    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)

# 对数据集进行预处理
tokenized_dataset = dataset.map(preprocess_function, batched=True)

# 4. 定义训练参数
training_args = TrainingArguments(
    output_dir="./results",      # 模型保存路径
    overwrite_output_dir=True,  # 如果目录已存在，覆盖
    evaluation_strategy="epoch",  # 每个 epoch 进行评估
    learning_rate=5e-5,          # 学习率
    per_device_train_batch_size=4,  # 每个设备的 batch size
    per_device_eval_batch_size=4,   # 每个设备的 eval batch size
    num_train_epochs=3,          # 训练 epoch 数
    save_strategy="epoch",       # 每个 epoch 保存一次模型
    save_total_limit=2,          # 最多保存的模型数量
    logging_dir="./logs",        # 日志保存路径
    logging_steps=100,           # 每 100 步记录一次日志
    load_best_model_at_end=True, # 保存最佳模型
    evaluation_strategy="epoch", # 每个 epoch 进行评估
    fp16=True,                   # 使用混合精度训练（需支持 GPU）
    report_to="none",            # 禁用默认的日志报告工具，如 WandB
)

# 5. 使用 Trainer 进行训练
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    tokenizer=tokenizer,
)

# 开始训练
trainer.train()

# 6. 保存模型
output_model_dir = "./fine_tuned_model"
os.makedirs(output_model_dir, exist_ok=True)
model.save_pretrained(output_model_dir)
tokenizer.save_pretrained(output_model_dir)

print(f"微调后的模型已保存到 {output_model_dir}")

代码说明

  1. 数据加载与预处理：
    • 数据集使用了 load_dataset 函数加载一个纯文本文件，每行是一条样本。
    • 在 preprocess_function 中，将文本数据转化为模型可以接受的输入格式（如 input_ids 和 attention_mask）。
  2. 模型选择：
    • 使用了 Hugging Face 的 GPT-2 模型（GPT2LMHeadModel）。你可以替换为其他支持的模型（如 GPT-Neo、Bloom、LLaMA）。
    • 如果你有特定领域的需求，可以选择更适合的预训练模型。
  3. 训练参数：
    • 超参数（如学习率、批次大小）可以根据实际情况调整。
    • 使用 fp16=True 开启混合精度训练，以减少显存占用并提升训练速度（需支持 GPU）。
  4. 模型保存：
    • 训练完成后，保存模型和分词器到指定目录，可以方便后续加载和部署。

文件结构建议

假设项目的文件结构如下：

project/
├── data/
│   └── train.txt       # 训练数据集，每行一条样本
├── fine_tuned_model/   # 微调后的模型保存路径
├── results/            # 训练过程的日志和检查点
└── train.py            # 上述训练代码

如何运行代码

  1. 将训练数据集（如 train.txt）存放在 data/ 文件夹下，每行一条文本。
  2. 运行代码文件 train.py：python train.py
  3. 微调完成后，模型会保存在 fine_tuned_model/ 目录下。

扩展功能

  1. 评估与测试：
    • 如果需要在验证集上评估模型，可以在 Trainer 中传入 eval_dataset 参数，并在 training_args 中启用 evaluation_strategy。
  2. 多卡训练：
    • 若有多块 GPU，可以利用 torch.distributed 实现分布式训练（Hugging Face 会自动检测多卡）。
    • 运行命令：torchrun --nproc_per_node=4 train.py 其中 --nproc_per_node=4 指定 GPU 数量。
  3. 迁移到云端：
    • 如果本地硬件不足，可以将代码迁移到云平台（如 AWS、Google Cloud、Azure）并利用其 GPU 实例。

适用场景

这段代码适用于以下场景：

  • 微调语言模型处理特定领域文本（如漏洞描述、攻击路径生成）。
  • 定制化模型生成特定输出（如漏洞修复建议）。
  • NLP 任务的快速实验（如分类、生成）。

如果有更多需求（如图数据或代码数据的处理），可以进一步扩展！

Print Friendly, PDF & Email