本文展示如何使用两种方式训练一个 GPT-2 Student 小模型一种是加载原始 GPT-2 模型直接蒸馏v1另一种是使用GPT2Config构建结构更小的模型进行压缩蒸馏v2。同时提供完整训练流程、代码与对比说明。 什么是模型蒸馏Distillation模型蒸馏是一种将大模型Teacher知识压缩到小模型Student的方法Teacher 输出 logits作为软标签指导 student 学习训练目标不再是标签而是“模仿”大模型的行为常用于模型压缩、推理加速、边缘部署 版本对比v1 vs v2方面distill_training.pyv1distill_training_v2.pyv2Student 构建方式直接加载 GPT2使用GPT2Config构建 6层小模型Transformer 层数126模型大小与 teacher 一样小一半场景适配蒸馏演示模型压缩 蒸馏 步骤通用结构两个版本都包含1️⃣ 加载模型和分词器teacher_modelGPT2LMHeadModel.from_pretrained(path_to_teacher).eval()tokenizerGPT2Tokenizer.from_pretrained(path_to_teacher)# v1加载标准 GPT2 作为 studentstudent_modelGPT2LMHeadModel.from_pretrained(gpt2).train()# v2构建 6层 transformer 小模型configGPT2Config(n_layer6,n_embd768,n_head12,vocab_size50257)student_modelGPT2LMHeadModel(config).train()vocab_size 必须一致否则 tokenizer 报错n_layer 越小模型越轻量蒸馏越有意义2️⃣ 构造数据集与 DataLoaderclassTextDataset(Dataset):def__init__(self,texts,tokenizer,max_length64):self.textstexts self.tokenizertokenizer self.max_lengthmax_lengthdef__getitem__(self,idx):inputsself.tokenizer(self.texts[idx],return_tensorspt,paddingmax_length,truncationTrue,max_lengthself.max_length)return{k:v.squeeze(0)fork,vininputs.items()}训练数据示例train_texts[Hello world!,The sky is blue.,AI is changing the world.,...]3️⃣ 蒸馏训练核心逻辑KL Lossloss_fntorch.nn.KLDivLoss(reductionbatchmean)lossloss_fn(student_logits.log_softmax(dim-1),teacher_logits.softmax(dim-1))使用 softmax log_softmax 组合符合 KL 散度要求不用交叉熵因为没有 ground truth label只有 teacher 输出作为 soft target4️⃣ 梯度更新流程optimizertorch.optim.AdamW(student_model.parameters(),lr5e-5)forbatchindataloader:input_idsbatch[input_ids].to(device)attention_maskbatch[attention_mask].to(device)withtorch.no_grad():teacher_logitsteacher_model(input_idsinput_ids,attention_maskattention_mask).logits student_logitsstudent_model(input_idsinput_ids,attention_maskattention_mask).logits lossloss_fn(student_logits.log_softmax(dim-1),teacher_logits.softmax(dim-1))optimizer.zero_grad()loss.backward()optimizer.step() 模型保存save_path./gpt2_student_v2# 或 gpt2_studentv1student_model.save_pretrained(save_path)tokenizer.save_pretrained(save_path)训练完成后即可用于推理支持使用 Hugging Face 接口加载。 总结蒸馏是一种让 student 模型模仿 teacher 的训练方式v1是结构不变的蒸馏偏向学习机制v2是结构压缩 蒸馏更适合部署使用 KL Loss 衡量两个模型输出概率分布的差异 本文为 GPT-2 蒸馏压缩项目第一篇共3篇第一篇从蒸馏到压缩两种方式训练GPT-2小模型 Student第二篇GPT-2 蒸馏模型推理实战标准 Student vs 压缩 Student 的调用对比第三篇GPT-2 蒸馏小模型部署实战Flask 封装推理接口与网页调用演示YoanAILab 技术导航页 项目源码 × 实战部署 × 转型经验一页总览 点击查看完整导航页 包含内容 GPT-2 项目源码GitHub✍️ CSDN 技术专栏合集 知乎转型日志