使用ptuning做意图分类
代码：
from torch.utils.data import Dataset, DataLoader, random_split
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import peft
import transformers
from peft import PrefixTuningConfig, get_peft_model, PeftModel
import time
from pathlib import Path
# ===========================
# 1. 环境与基础配置
# ===========================
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"PyTorch版本: {torch.__version__}")
print(f"设备: {device}")
print(f"CUDA版本: {torch.version.cuda}")
print(f"Transformers版本: {transformers.__version__}")
print(f"PEFT版本: {peft.__version__}")
# 基础模型名称（本地路径或模型名）
# base_model_name = "Qwen/Qwen3-0.6B"
base_model_name = r"D:\source\models\Qwen3-0.6B"
# Prefix Tuning 权重保存目录
ptuning_output_dir = "ptuning_intent_qwen3_0.6b"
# ===========================
# 2. 示例数据
# ===========================
# ⚠️ 实际使用时，一定要扩充更多样本，这里只是演示
sample_data = [
{
"instruction": "作为意图分类专家，只能输出以下两个标签之一：【攻击研判】或【攻击分析】。不要输出其他内容。",
"input": "研判告警 9d34dc99-a981-41dc-94bf-10b282354da4",
"output": "攻击研判"
},
{
"instruction": "作为意图分类专家，只能输出以下两个标签之一：【攻击研判】或【攻击分析】。不要输出其他内容。",
"input": "列举10条高危告警",
"output": "攻击分析"
}
]
# ===========================
# 3. 构建模型 + Prefix Tuning (P-Tuning)
# ===========================
def setup_model(model_name):
"""
加载基础模型和 tokenizer，并应用 Prefix Tuning（P-Tuning 风格）。
当前环境为 CPU：不使用 device_map='auto'，避免 meta tensor / offload。
"""
tokenizer = AutoTokenizer.from_pretrained(
model_name,
local_files_only=True,
trust_remote_code=True
)
# 加载基础模型
model = AutoModelForCausalLM.from_pretrained(
model_name,
local_files_only=True,
dtype=torch.float32, # 新版 transformers 建议用 dtype
device_map=None, # 不启用 auto/offload，避免 meta tensor
trust_remote_code=True
)
"""
在你当前的 Prefix Tuning 代码里：
python
prefix_config = PrefixTuningConfig(
task_type="CAUSAL_LM",
num_virtual_tokens=32,
)
model = get_peft_model(model, prefix_config)
PEFT 会帮你做的事情大致是：
给每一层追加长度为 32 的虚拟前缀（可以理解为 32 个“隐形 token”）
只训练这些前缀相关的参数
在你训练时，它会学出一种“一旦看到这个前缀，后面的输入就应该被理解为某种特定分类任务”的内在模式
推理时，只要你加载了这个 prefix，模型就会在这个任务风格下工作
"""
# Prefix Tuning 配置（P-Tuning 思路）
prefix_config = PrefixTuningConfig(
task_type="CAUSAL_LM",
num_virtual_tokens=32, # 虚拟 prompt 长度，可根据需要调整：16/32/64
# 如果报 encoder_hidden_size 相关错误，再启用下面这一行：
# encoder_hidden_size=model.config.hidden_size
)
# 应用 Prefix Tuning
model = get_peft_model(model, prefix_config)
# 打印可训练参数
print("=== 模型参数统计 ===")
model.print_trainable_parameters()
return model, tokenizer
# ===========================
# 4. 数据集定义
# ===========================
class IntentClassificationDataset(Dataset):
def __init__(self, data, tokenizer, max_length=256):
self.data = data
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
item = self.data[idx]
# 构建完整序列
full_text = (
f"指令：{item['instruction']}\n"
f"输入：{item['input']}\n"
f"输出：{item['output']}{self.tokenizer.eos_token}"
)
instruction_text = (
f"指令：{item['instruction']}\n"
f"输入：{item['input']}\n"
f"输出："
)
# Tokenize 完整序列
encoding = self.tokenizer(
full_text,
max_length=self.max_length,
padding='max_length',
truncation=True,
return_tensors='pt'
)
# Tokenize 指令部分，计算需要 mask 的长度
instruction_encoding = self.tokenizer(
instruction_text,
return_tensors='pt'
)
instruction_length = instruction_encoding['input_ids'].shape[1]
# 创建 labels 并 mask 掉指令部分
labels = encoding['input_ids'].clone()
labels[0, :instruction_length] = -100 # -100 会在 loss 里被忽略
return {
'input_ids': encoding['input_ids'].squeeze(0),
'attention_mask': encoding['attention_mask'].squeeze(0),
'labels': labels.squeeze(0)
}
# ===========================
# 5. 训练 + 保存 Prefix Tuning 权重
# ===========================
def train_and_save():
model, tokenizer = setup_model(base_model_name)
# 不再调用 model.to(device)，setup_model 已经在 CPU 上加载好了
dataset = IntentClassificationDataset(sample_data, tokenizer, max_length=256)
# 简单划分训练/验证
train_size = int(0.8 * len(dataset))
eval_size = len(dataset) - train_size
train_dataset, eval_dataset = random_split(dataset, [train_size, eval_size])
print(f"训练集大小: {len(train_dataset)}")
print(f"验证集大小: {len(eval_dataset)}")
batch_size = 1 # CPU 训练，数据也少，小 batch 就行
train_dataloader = DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=0
)
eval_dataloader = DataLoader(
eval_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=0
)
learning_rate = 5e-5
weight_decay = 0.01
num_epochs = 3 # 比之前多训几轮
torch.manual_seed(123)
optimizer = torch.optim.AdamW(
model.parameters(),
lr=learning_rate,
weight_decay=weight_decay
)
start_time = time.time()
global_step = 0
for epoch in range(num_epochs):
model.train()
for batch in train_dataloader:
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['labels'].to(device)
optimizer.zero_grad()
outputs = model(
input_ids=input_ids,
attention_mask=attention_mask,
labels=labels
)
loss = outputs.loss
loss.backward()
optimizer.step()
global_step += 1
print(f"Epoch {epoch}, step {global_step}, loss {loss.item():.4f}")
# 可选的简单验证
model.eval()
if len(eval_dataset) > 0:
eval_losses = []
with torch.no_grad():
for batch in eval_dataloader:
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['labels'].to(device)
outputs = model(
input_ids=input_ids,
attention_mask=attention_mask,
labels=labels
)
eval_losses.append(outputs.loss.item())
print(f"Epoch {epoch} 验证集平均loss: {sum(eval_losses)/len(eval_losses):.4f}")
end_time = time.time()
execution_time_minutes = (end_time - start_time) / 60
print(f"Training completed in {execution_time_minutes:.2f} minutes.")
# ====== 保存 Prefix Tuning 增量权重（不会保存基础模型参数）======
Path(ptuning_output_dir).mkdir(parents=True, exist_ok=True)
model.save_pretrained(ptuning_output_dir)
tokenizer.save_pretrained(ptuning_output_dir)
print(f"P-Tuning 权重与 tokenizer 已保存到: {ptuning_output_dir}")
# ===========================
# 6. 加载 Prefix Tuning + 推理示例
# ===========================
def load_ptuning_and_infer(user_input: str):
"""
读取已训练好的 Prefix Tuning 权重，进行推理。
"""
# 1. 加载基础模型和 tokenizer（同样全部放 CPU）
tokenizer = AutoTokenizer.from_pretrained(
base_model_name,
local_files_only=True,
trust_remote_code=True
)
base_model = AutoModelForCausalLM.from_pretrained(
base_model_name,
local_files_only=True,
dtype=torch.float32,
device_map=None,
trust_remote_code=True
)
# 2. 将 Prefix Tuning 适配器加载到基础模型上（也在 CPU）
model = PeftModel.from_pretrained(
base_model,
ptuning_output_dir,
device_map=None
)
model.eval()
# 3. 构造推理提示词（和训练时格式保持一致）
instruction = "作为意图分类专家，只能输出以下两个标签之一：【攻击研判】或【攻击分析】。不要输出其他内容。"
prompt = f"指令：{instruction}\n输入：{user_input}\n输出："
inputs = tokenizer(
prompt,
return_tensors="pt"
).to(device)
with torch.no_grad():
generated_ids = model.generate(
**inputs,
max_new_tokens=128, # 分类任务只需要很短的输出
do_sample=False, # 分类任务一般不需要采样
pad_token_id=tokenizer.eos_token_id
)
# 只取新生成的部分
generated_text = tokenizer.decode(
generated_ids[0][inputs['input_ids'].shape[1]:],
skip_special_tokens=True
)
text = generated_text.strip()
# 简单后处理，只认两个标签
if "攻击研判" in text:
intent = "攻击研判"
elif "攻击分析" in text:
intent = "攻击分析"
else:
# 兜底策略（你可以改成 "未知"）
intent = "攻击分析"
print(f"模型预测意图: {intent}（原始输出: {text}）")
return intent
# ===========================
# 7. 主入口
# ===========================
if __name__ == "__main__":
# 第一步：训练并保存 Prefix Tuning
train_and_save()
# 第二步：加载 Prefix Tuning 并推理
print("\n==== 推理示例 ====")
test_inputs = [
"研判告警 9d34dc99-a981-41dc-94bf-10b282354da4",
"请列举10条高危告警"
]
for t in test_inputs:
print(f"\n用户输入: {t}")
_ = load_ptuning_and_infer(t)
1. 大局观：P-Tuning 在干嘛？
一句话概括：
P-Tuning 不是去改模型里的权重，而是学一段“虚拟的提示词”嵌入向量，让模型在这段提示的引导下，更好完成下游任务。
也就是说：
- 模型主体参数（几亿、几十亿的权重）通常是冻结的
- 只新增一小块参数：一段长度为 LL 的「可学习的 tokens 嵌入」
- 训练时，只更新这段嵌入参数
- 推理时，把这段嵌入“拼接”到输入前面，一起喂给模型
和“往前面加一串文字 prompt”很像，但区别是：
- 普通 prompt：是人类写的可见文字，通过 tokenizer 变成 token，再查 embedding
- P-Tuning 的 prompt：是直接在 embedding 空间里学出来的“虚拟 token”，不一定对应任何真实词
2. 三个相关概念：Prompt Tuning、P-Tuning、Prefix Tuning
很多文章会混着说，简单对一下：
对你这种应用：
- 基础模型：Qwen3-0.6B 已经学到大量语言和安全相关的知识
- 你的任务：只是教它“看输入句子 -> 输出两种标签之一”
- 不需要重塑模型，只需要“告诉它该怎么用已有能力来做分类”
- P-Tuning / LoRA 这类参数高效微调就非常合适
6. 结合你现在代码，可以脑补一下“内部图像”
在你当前的 Prefix Tuning 代码里：
python
PEFT 会帮你做的事情大致是：
- 给每一层追加长度为 32 的虚拟前缀（可以理解为 32 个“隐形 token”）
- 只训练这些前缀相关的参数
- 在你训练时，它会学出一种“一旦看到这个前缀，后面的输入就应该被理解为某种特定分类任务”的内在模式
- 推理时，只要你加载了这个 prefix，模型就会在这个任务风格下工作
附：使用loar微调的方式
from torch.utils.data import Dataset, DataLoader, random_split
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import peft
import transformers
from peft import LoraConfig, get_peft_model, PeftModel
import time
from pathlib import Path
# ===========================
# 1. 环境与基础配置
# ===========================
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"PyTorch版本: {torch.__version__}")
print(f"设备: {device}")
print(f"CUDA版本: {torch.version.cuda}")
print(f"Transformers版本: {transformers.__version__}")
print(f"PEFT版本: {peft.__version__}")
# 基础模型名称
# base_model_name = "Qwen/Qwen3-0.6B" # 如果想换8B就改这里
base_model_name = r"D:\source\models\Qwen3-0.6B"
# LoRA 权重保存目录
lora_output_dir = "lora_intent_qwen3_0.6b"
# ===========================
# 2. 示例数据
# ===========================
sample_data = [
{
"instruction": "作为意图分类专家，分析用户输入，将输入分类为攻击研判或攻击分析",
"input": "研判告警 9d34dc99-a981-41dc-94bf-10b282354da4",
"output": "攻击研判"
},
{
"instruction": "作为意图分类专家，分析用户输入，将输入分类为攻击研判或攻击分析",
"input": "列举10条高危告警",
"output": "攻击分析"
}
]
# ===========================
# 3. 构建模型 + LoRA
# ===========================
def setup_model(model_name):
# 本地加载，如需联网加载可去掉 local_files_only=True
tokenizer = AutoTokenizer.from_pretrained(
model_name,
local_files_only=True,
trust_remote_code=True
)
model = AutoModelForCausalLM.from_pretrained(
model_name,
local_files_only=True,
dtype="auto",
device_map="auto",
trust_remote_code=True
)
# LoRA配置
lora_config = LoraConfig(
r=8,
lora_alpha=32,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], # Qwen3 相关模块
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM"
)
# 应用LoRA
model = get_peft_model(model, lora_config)
# 打印可训练参数
print("=== 模型参数统计 ===")
model.print_trainable_parameters()
return model, tokenizer
# ===========================
# 4. 数据集定义
# ===========================
class IntentClassificationDataset(Dataset):
def __init__(self, data, tokenizer, max_length=256):
self.data = data
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
item = self.data[idx]
# 构建完整序列
full_text = (
f"指令：{item['instruction']}\n"
f"输入：{item['input']}\n"
f"输出：{item['output']}{self.tokenizer.eos_token}"
)
instruction_text = (
f"指令：{item['instruction']}\n"
f"输入：{item['input']}\n"
f"输出："
)
# Tokenize 完整序列
encoding = self.tokenizer(
full_text,
max_length=self.max_length,
padding='max_length',
truncation=True,
return_tensors='pt'
)
# Tokenize 指令部分，计算需要 mask 的长度
instruction_encoding = self.tokenizer(
instruction_text,
return_tensors='pt'
)
instruction_length = instruction_encoding['input_ids'].shape[1]
# 创建 labels 并 mask 掉指令部分
labels = encoding['input_ids'].clone()
labels[0, :instruction_length] = -100 # -100 会在 loss 里被忽略
return {
'input_ids': encoding['input_ids'].squeeze(0),
'attention_mask': encoding['attention_mask'].squeeze(0),
'labels': labels.squeeze(0)
}
# ===========================
# 5. 训练 + 保存 LoRA
# ===========================
def train_and_save():
model, tokenizer = setup_model(base_model_name)
model.to(device)
dataset = IntentClassificationDataset(sample_data, tokenizer, max_length=256)
# 简单划分训练/验证
train_size = int(0.8 * len(dataset))
eval_size = len(dataset) - train_size
train_dataset, eval_dataset = random_split(dataset, [train_size, eval_size])
print(f"训练集大小: {len(train_dataset)}")
print(f"验证集大小: {len(eval_dataset)}")
batch_size = 1 # 显存紧张时小一点
train_dataloader = DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=0
)
eval_dataloader = DataLoader(
eval_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=0
)
learning_rate = 5e-5
weight_decay = 0.01
num_epochs = 1
torch.manual_seed(123)
optimizer = torch.optim.AdamW(
model.parameters(),
lr=learning_rate,
weight_decay=weight_decay
)
start_time = time.time()
global_step = 0
for epoch in range(num_epochs):
model.train()
for batch in train_dataloader:
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['labels'].to(device)
optimizer.zero_grad()
outputs = model(
input_ids=input_ids,
attention_mask=attention_mask,
labels=labels
)
loss = outputs.loss
loss.backward()
optimizer.step()
global_step += 1
if global_step % 1 == 0:
print(f"Epoch {epoch}, step {global_step}, loss {loss.item():.4f}")
end_time = time.time()
execution_time_minutes = (end_time - start_time) / 60
print(f"Training completed in {execution_time_minutes:.2f} minutes.")
# ====== 保存 LoRA 增量权重（不会保存基础模型参数）======
Path(lora_output_dir).mkdir(parents=True, exist_ok=True)
model.save_pretrained(lora_output_dir)
tokenizer.save_pretrained(lora_output_dir)
print(f"LoRA 权重与 tokenizer 已保存到: {lora_output_dir}")
# ===========================
# 6. 加载 LoRA + 推理示例
# ===========================
def load_lora_and_infer(user_input: str):
"""
读取已训练好的 LoRA 权重，进行推理。
这里假定 LoRA 已经保存在 lora_output_dir。
"""
# 1. 加载基础模型和 tokenizer
tokenizer = AutoTokenizer.from_pretrained(
base_model_name,
local_files_only=True,
trust_remote_code=True
)
base_model = AutoModelForCausalLM.from_pretrained(
base_model_name,
local_files_only=True,
dtype="auto",
device_map="auto",
trust_remote_code=True
)
# 2. 将 LoRA 适配器加载到基础模型上
model = PeftModel.from_pretrained(
base_model,
lora_output_dir,
device_map="auto"
)
model.eval()
# 3. 构造推理提示词（和训练时格式保持一致）
instruction = "作为意图分类专家，分析用户输入，将输入分类为攻击研判或攻击分析"
prompt = f"指令：{instruction}\n输入：{user_input}\n输出："
inputs = tokenizer(
prompt,
return_tensors="pt"
).to(model.device)
with torch.no_grad():
generated_ids = model.generate(
**inputs,
max_new_tokens=128,
do_sample=False, # 分类任务一般不需要采样
pad_token_id=tokenizer.eos_token_id
)
# 只取新生成的部分
generated_text = tokenizer.decode(
generated_ids[0][inputs['input_ids'].shape[1]:],
skip_special_tokens=True
)
# 简单清洗一下空格和换行
intent = generated_text.strip()
print(f"模型预测意图: {intent}")
return intent
# ===========================
# 7. 主入口
# ===========================
if __name__ == "__main__":
# 第一步：训练并保存 LoRA
train_and_save()
# 第二步：加载 LoRA 并推理
print("\n==== 推理示例 ====")
test_inputs = [
"研判告警 9d34dc99-a981-41dc-94bf-10b282354da4",
"请列举10条高危告警"
]
for t in test_inputs:
print(f"\n用户输入: {t}")
_ = load_lora_and_infer(t)
浙公网安备 33010602011771号