下面我将介绍如何使用Hugging Face的Transformer框架对微调后的Qwen或DeepSeek模型进行非流式批量推理。
一、准备工作
首先确保已安装必要的库：
pip install transformers torch二、批量推理实现代码
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from typing import List
class BatchInference:
def __init__(self, model_path: str, device: str = "cuda" if torch.cuda.is_available() else "cpu"):
"""
初始化模型和tokenizer
:param model_path: 微调后的模型路径(Hugging Face模型ID或本地路径)
:param device: 推理设备(cpu/cuda)
"""
self.device = device
self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
self.model = AutoModelForCausalLM.from_pretrained(
model_path,
trust_remote_code=True,
torch_dtype=torch.float16 if "cuda" in device else torch.float32,
device_map="auto"
)
self.model.eval()
# 设置pad_token_id为eos_token_id如果不存在pad_token
if self.tokenizer.pad_token_id is None:
self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
def batch_predict(self, prompts: List[str], max_new_tokens: int = 512, batch_size: int = 4, **generate_kwargs) -> List[str]:
"""
批量推理方法
:param prompts: 输入提示列表
:param max_new_tokens: 生成的最大token数
:param batch_size: 批量大小
:param generate_kwargs: 额外的生成参数
:return: 生成结果列表
"""
all_results = []
# 分批次处理
for i in range(0, len(prompts), batch_size):
batch_prompts = prompts[i:i + batch_size]
# 编码输入
inputs = self.tokenizer(
batch_prompts,
padding=True,
truncation=True,
return_tensors="pt",
max_length=1024 # 可根据需要调整
).to(self.device)
# 生成输出
with torch.no_grad():
outputs = self.model.generate(
**inputs,
max_new_tokens=max_new_tokens,
pad_token_id=self.tokenizer.pad_token_id,
**generate_kwargs
)
# 解码输出并移除输入部分
batch_results = []
for j in range(len(outputs)):
output = outputs[j]
input_length = inputs["input_ids"][j].shape[0]
generated = output[input_length:]
batch_results.append(self.tokenizer.decode(generated, skip_special_tokens=True))
all_results.extend(batch_results)
return all_results
# 使用示例
if __name__ == "__main__":
# 替换为你的模型路径(本地或Hugging Face模型ID)
MODEL_PATH = "Qwen/Qwen-7B-Chat" # 或"deepseek-ai/deepseek-llm-7b"等
# 初始化推理器
inferencer = BatchInference(MODEL_PATH)
# 示例输入
prompts = [
"请解释一下量子计算的基本原理",
"写一首关于春天的诗",
"如何用Python实现快速排序?",
"Transformer模型的主要创新点是什么?"
]
# 批量推理
results = inferencer.batch_predict(
prompts,
max_new_tokens=256,
batch_size=2, # 根据GPU内存调整
temperature=0.7,
top_p=0.9
)
# 打印结果
for prompt, result in zip(prompts, results):
print(f"输入: {prompt}\n输出: {result}\n{'-'*50}")三、关键点说明
-
设备管理:
-
自动检测并使用可用的GPU
-
支持半精度(fp16)推理以节省显存
-
-
批量处理:
-
将输入分成小批次处理，避免内存不足
-
自动填充(padding)使批次内样本长度一致
-
-
生成参数:
-
max_new_tokens: 控制生成的最大长度 -
temperature和top_p: 控制生成的随机性 -
可通过
generate_kwargs传递其他生成参数
-
-
内存优化:
-
使用
torch.no_grad()减少内存消耗 -
根据GPU内存调整
batch_size
-
四、进阶优化
-
使用Flash Attention (如果模型支持):
model = AutoModelForCausalLM.from_pretrained(
model_path,
trust_remote_code=True,
torch_dtype=torch.float16,
use_flash_attention_2=True, # 启用Flash Attention
device_map="auto"
)
2.量化推理 (减少显存使用):
model = AutoModelForCausalLM.from_pretrained(
model_path,
trust_remote_code=True,
torch_dtype=torch.float16,
load_in_8bit=True, # 8位量化
device_map="auto"
)3.使用vLLM等优化库 (对于生产环境):
from vllm import LLM, SamplingParams
llm = LLM(model=MODEL_PATH)
sampling_params = SamplingParams(temperature=0.7, top_p=0.9, max_tokens=256)
outputs = llm.generate(prompts, sampling_params)以上代码提供了基于Transformer框架的Qwen/DeepSeek模型批量推理基础实现，可根据实际需求进行调整和优化。