Hugging Face Transformers全模块串联实战:BERT中文分类深度解析

发布时间:2026/7/12 11:34:39

Hugging Face Transformers全模块串联实战:BERT中文分类深度解析 1. 项目概述为什么这个BERT案例是Transformers框架的“终极考卷”你有没有试过把Hugging Face Transformers框架从头到尾串一遍不是只跑通一个Trainer.train()而是真正搞懂模型怎么加载、数据怎么喂、损失怎么算、梯度怎么传、日志怎么打、checkpoint怎么存、评估怎么验、预测怎么出——所有模块像齿轮一样咬合转动缺一不可。这个标题里的“17-Hugging Face Transformers之BERT 案例详解Transformers 框架全模块串联”说白了就是一份框架级操作手册它不教你怎么调参而是带你亲手把Transformers这台精密仪器的每一个螺丝、每一根管线、每一块电路板都拆开、看清、再装回去。我做NLP工程落地六年带过二十多个工业级文本分类项目最常被问的问题不是“BERT怎么微调”而是“为什么我的Trainer报错no module named transformers”、“为什么eval时loss突然飙升”、“为什么保存的模型load不回来”——这些问题90%都源于对框架内部协作逻辑的模糊认知。这个案例就是专治这种“会用但不懂”的典型症状。它选BERT是因为BERT结构干净无decoder、任务明确分类、生态成熟THUCNews等中文数据集现成是验证框架完整性的最佳“探针”。而“全模块串联”四个字意味着你要看到AutoModel如何根据配置自动分发到BertModelDataCollatorWithPadding如何与tokenizer协同处理变长序列Trainer如何在_inner_training_loop里调度CallbackHandler、OptimizerManager和EvalLoopContainer甚至JITCheckpointCallback如何在SIGTERM信号到来前3秒抢出一个完整checkpoint。这不是一个“能跑就行”的玩具Demo而是一次对Transformers设计哲学的深度测绘约定优于配置、策略模式解耦、观察者模式驱动、装饰器增强鲁棒性。如果你的目标是成为能独立搭建、调试、优化任何基于Transformers的NLP流水线的工程师而不是只会复制粘贴pipeline调用的调包侠那么这个案例就是你绕不开的必经之路。2. 核心模块串联逻辑从BERT分类任务反向推导框架设计2.1 为什么是BERT分类——任务驱动的模块选择逻辑我们先锁定场景基于BERT对THUCNews新闻标题做多分类。这是一个典型的监督学习下游任务输入是短文本标题输出是离散标签如“体育”、“财经”、“娱乐”。这个看似简单的任务恰恰是检验Transformers框架完整性的黄金标尺原因有三第一它强制暴露所有关键路径。预训练模型BertModel必须加载文本需要分词BertTokenizer标签需编码LabelEncoder数据要组织成batchDataCollatorWithPadding训练需迭代Trainer评估需指标compute_metrics预测需接口Trainer.predict。任何一个环节缺失或错位整个流程就断掉。第二它天然规避了decoder的复杂性。GPT类模型涉及generate()、decoder_input_ids、attention_mask的双向交互而BERT作为encoder-only模型其forward()逻辑清晰输入input_ids和attention_mask输出last_hidden_state再接一个nn.Linear层做分类。这让我们能把全部注意力聚焦在框架的“骨架”上而非被生成逻辑的“血肉”干扰。第三它完美适配中文实践痛点。THUCNews是中文NLP领域最常用的基准数据集之一但它的原始格式纯文本目录结构与Transformers要求的Dataset对象之间存在鸿沟。这个案例必须解决如何用load_dataset(thucnews)或手动构建Dataset.from_dict()如何处理中文分词特有的[UNK]率问题如何为DataCollator指定正确的padding_sideright因为BERT的[CLS]在开头padding必须在末尾如何确保label2id映射在Trainer初始化时被正确注入到模型的config.id2label中。这些都不是BERT模型本身的问题而是Transformers框架如何将模型、数据、训练器三者无缝缝合的工程细节。所以这个案例的起点不是“我要用BERT”而是“我要让Transformers框架为我服务”BERT只是那个最听话、最透明的执行者。2.2 全模块串联的四层架构从基础设施到业务逻辑当你运行trainer.train()时背后并非一个黑箱而是一个精密的四层架构在协同工作。理解这四层是读懂“全模块串联”的钥匙。第一层是基础设施层这是整个框架的地基。它包含TrainingArguments——一个dataclass它把所有超参数per_device_train_batch_size16,learning_rate2e-5,num_train_epochs3打包成一个可序列化、可命令行解析的对象。更重要的是它的cached_property机制如device,n_gpu,parallel_mode实现了“按需计算”避免了初始化时的冗余探测。比如device属性会自动按CUDA NPU MPS CPU的优先级查找可用设备你无需写if torch.cuda.is_available(): ...。第二层是配置与参数层它负责将抽象的配置翻译成具体的对象。AutoModel和AutoTokenizer是这一层的明星它们根据model_name_or_pathbert-base-chinese自动匹配并实例化BertModel和BertTokenizer省去了你手动导入from transformers import BertModel, BertTokenizer的麻烦。DataCollatorForSequenceClassification则根据任务类型这里是分类自动选择DataCollatorWithPadding并预设好label_pad_token_id-100因为分类任务的labels是标量不需要padding但框架统一用-100忽略loss计算。第三层是核心训练层即Trainer本体。它不是一个万能胶而是一个高度解耦的协调中心。它的train()方法只是一个入口真正的核心是_inner_training_loop()一个被find_executable_batch_size装饰器包裹的函数。这个装饰器是框架鲁棒性的体现当你的GPU显存不足OOM时它不会直接崩溃而是自动将per_device_train_batch_size乘以0.9重试直到找到一个能跑通的尺寸。第四层是扩展与增强层它赋予框架无限可能。CallbackSystem是这一层的灵魂EarlyStoppingCallback、ProgressCallback、JITCheckpointCallback都是插件你可以随时增删。compute_loss_func允许你替换默认的交叉熵损失比如加入Focal Loss解决类别不平衡compute_metrics让你自定义评估指标不只是accuracy还可以是precision、recall、f1-score的加权平均。这四层不是线性堆叠而是网状依赖Trainer初始化时TrainingArguments提供参数AutoModel提供模型DataCollator提供数据整理器Callback提供行为钩子。任何一个模块的变更都会通过依赖链影响全局。所谓“串联”就是看清这张网并亲手把它织起来。2.3 “串联”的本质模块间的数据流与控制流“全模块串联”的“串联”二字绝非物理上的连接而是指数据流Data Flow和控制流Control Flow在模块间的精确传递与响应。我们以一个训练step为例追踪一次完整的“心跳”。首先数据流启动Trainer.get_train_dataloader()被调用它内部使用LengthGroupedSampler来自trainer_pt_utils对数据集索引进行排序将长度相近的样本分到同一batch极大减少padding带来的显存浪费。接着DataCollatorWithPadding接手它接收一个list[dict]每个dict是{input_ids: [...], attention_mask: [...], label: 0}将其pad成统一长度的{input_ids: [batch_size, max_len], attention_mask: [batch_size, max_len], labels: [batch_size]}张量字典。这个字典被送入training_step()。此时控制流介入training_step()先调用self._prepare_inputs(inputs)将inputs移动到args.device上然后进入compute_loss()这里触发三层优先级如果用户设置了compute_loss_func就用它否则看是否有label_smoother最后才取模型输出的outputs[loss]。compute_loss()返回的loss是一个标量张量它被self.accelerator.backward(loss)反向传播梯度累积在模型参数上。当gradient_accumulation_steps满足时optimizer.step()被执行lr_scheduler.step()随之更新学习率。与此同时控制流并未停止callback_handler.on_step_end()被触发DefaultFlowCallback检查state.global_step % args.logging_steps 0若为真则设置control.should_log TrueEarlyStoppingCallback则检查metrics[eval_f1]是否比历史最佳值更好若连续patience3次未提升则设置control.should_training_stop True。最终control对象被返回Trainer根据其中的信号决定下一步动作。整个过程数据张量在DataCollator→Trainer→Model→Loss→Optimizer之间流动而控制信号则在Trainer→CallbackHandler→Callbacks→Trainer之间循环。这种数据与控制的分离正是Transformers框架高内聚、低耦合的设计精髓。你无法只改DataCollator而不影响Trainer的行为也无法只动Callback而不改变训练流程。它们被牢牢“串联”在同一个事件驱动的生命周期里。3. 实操步骤详解从零开始构建一个可复现的BERT分类流水线3.1 环境准备与依赖安装避开“no module named transformers”陷阱“no module named transformers”是新手遇到的第一个也是最经典的拦路虎。它看似简单实则暗藏玄机。我踩过的坑告诉我这个问题90%不是没装而是装错了地方或版本冲突。第一步绝对不要用pip install transformers在全局环境。你应该创建一个纯净的虚拟环境这是所有专业项目的起点。在终端执行python -m venv bert_env source bert_env/bin/activate # Linux/Mac # bert_env\Scripts\activate.bat # Windows激活后which python应指向bert_env/bin/python。第二步安装Transformers及其生态。官方推荐的安装方式是pip install transformers[torch,sentencepiece]这会同时安装PyTorch和SentencePiece中文分词必需。但请注意[torch,sentencepiece]是extras不是版本号别写成transformers[torch]4.40.0。如果你在国内pip源慢可以临时换清华源pip install -i https://pypi.tuna.tsinghua.edu.cn/simple/ transformers[torch,sentencepiece]。第三步验证安装是否成功且版本兼容。运行以下Python代码import transformers import torch print(fTransformers version: {transformers.__version__}) print(fPyTorch version: {torch.__version__}) print(fCUDA available: {torch.cuda.is_available()})关键点在于版本匹配。截至2024年transformers4.40.0与torch2.2.0是经过充分测试的稳定组合。如果你强行安装transformers4.45.0而PyTorch还是2.1.0就可能出现ImportError: cannot import name XXX from transformers。第四步处理常见冲突。如果你之前装过sentence-transformers它会自带一个旧版transformers导致冲突。解决方案是先pip uninstall sentence-transformers transformers再pip install transformers[torch,sentencepiece]最后pip install sentence-transformers它会自动适配已安装的transformers版本。第五步IDE配置。如果你用VS Code务必在设置中指定Python解释器为bert_env/bin/python否则编辑器里不报错但终端运行时报ModuleNotFoundError。PyCharm同理在Project Interpreter里选择bert_env的Python。这一步做完import transformers就不会再报错了。记住环境问题是地基地基不牢后面所有代码都是空中楼阁。我曾为一个客户排查三天最后发现是Jupyter Notebook kernel没切换到新环境这种低级错误值得花十分钟彻底杜绝。3.2 数据准备与预处理THUCNews的标准化处理流程THUCNews数据集是中文NLP的“Hello World”但它原始的文件结构train/体育/1.txt,train/财经/2.txt与Transformers要求的Dataset对象之间隔着一道需要亲手搭建的桥。我们采用最通用、最可控的手动构建法而非依赖load_dataset(thucnews)后者有时会因网络或缓存问题失败。第一步下载并解压数据。从清华大学NLP组官网下载THUCNews.zip解压到项目目录下的data/thucnews。第二步构建Dataset对象。我们不用pandas而是用datasets库的原生API保证内存效率from datasets import Dataset, Features, Value, ClassLabel import os import random def load_thucnews_data(data_dir, splittrain): 加载THUCNews数据返回Dataset对象 texts, labels [], [] # THUCNews的10个类别 categories [体育, 娱乐, 家居, 房产, 教育, 时尚, 时政, 游戏, 科技, 财经] label2id {cat: i for i, cat in enumerate(categories)} split_dir os.path.join(data_dir, split) for category in categories: cat_dir os.path.join(split_dir, category) if not os.path.exists(cat_dir): continue for file_name in os.listdir(cat_dir): if not file_name.endswith(.txt): continue file_path os.path.join(cat_dir, file_name) try: with open(file_path, r, encodingutf-8) as f: content f.read().strip() # THUCNews的标题在第一行正文在后面我们只取标题 title content.split(\n)[0].strip() if len(title) 2: # 过滤过短标题 continue texts.append(title) labels.append(label2id[category]) except Exception as e: print(fError reading {file_path}: {e}) continue # 构建Features明确数据类型 features Features({ text: Value(string), label: ClassLabel(namescategories) }) return Dataset.from_dict({text: texts, label: labels}, featuresfeatures) # 加载训练集和验证集 train_dataset load_thucnews_data(data/thucnews, train) val_dataset load_thucnews_data(data/thucnews, dev) print(fTrain samples: {len(train_dataset)}, Val samples: {len(val_dataset)})这段代码的关键在于Features的显式声明。ClassLabel不仅存储了类别名还自动创建了label2id和id2label映射这是后续Trainer能正确处理分类任务的基础。第三步分词与编码。使用BertTokenizer并注意中文的特殊性from transformers import BertTokenizer tokenizer BertTokenizer.from_pretrained(bert-base-chinese) def tokenize_function(examples): # 对text字段进行分词返回input_ids和attention_mask # truncationTrue确保超长文本被截断max_length128是BERT的常用长度 return tokenizer( examples[text], truncationTrue, paddingTrue, # 动态paddingbatch内最长为准 max_length128, return_tensorspt # 返回PyTorch张量 ) # 批量处理num_proc4利用多核CPU加速 tokenized_train train_dataset.map(tokenize_function, batchedTrue, num_proc4) tokenized_val val_dataset.map(tokenize_function, batchedTrue, num_proc4) # 移除原始text列只保留模型需要的input_ids等 tokenized_train tokenized_train.remove_columns([text]) tokenized_val tokenized_val.remove_columns([text]) # 将label列重命名为labels这是Trainer的约定 tokenized_train tokenized_train.rename_column(label, labels) tokenized_val tokenized_val.rename_column(label, labels)这里paddingTrue是关键它让DataCollator的工作变得简单。truncationTrue防止OOMmax_length128是经验之选THUCNews标题普遍较短128足够覆盖99%的样本且比512节省大量显存。第四步数据整理器DataCollator。我们不手动写而是用框架提供的DataCollatorWithPaddingfrom transformers import DataCollatorWithPadding # tokenizer是必需的它知道如何pad data_collator DataCollatorWithPadding(tokenizertokenizer, paddinglongest)paddinglongest表示在一个batch内所有样本pad到该batch中最长的那个长度这是最省内存的策略。至此数据准备完成tokenized_train是一个标准的Dataset对象其每个样本都是{input_ids: [...], attention_mask: [...], labels: 0}完全符合Trainer的输入规范。3.3 模型构建与配置从预训练权重到下游任务头构建BERT模型核心在于理解PreTrainedModel的继承体系。BertModel是基础编码器它只输出last_hidden_state而BertForSequenceClassification则在其之上添加了一个nn.Linear分类头。对于THUCNews的10分类任务我们必须用后者。第一步加载预训练模型与配置from transformers import BertForSequenceClassification, TrainingArguments, Trainer # 加载预训练的bert-base-chinese权重 model BertForSequenceClassification.from_pretrained( bert-base-chinese, num_labels10, # THUCNews有10个类别 id2label{i: cat for i, cat in enumerate([体育, 娱乐, 家居, 房产, 教育, 时尚, 时政, 游戏, 科技, 财经])}, label2id{cat: i for i, cat in enumerate([体育, 娱乐, 家居, 房产, 教育, 时尚, 时政, 游戏, 科技, 财经])} )from_pretrained()是魔法所在。它会自动下载bert-base-chinese的pytorch_model.bin模型权重和config.json模型结构并将权重加载到BertForSequenceClassification的对应层上。num_labels10告诉模型分类头的输出维度是10。id2label和label2id的显式传入至关重要它确保了模型内部的config与外部数据集的ClassLabel完全一致避免了预测时label映射错乱的灾难性错误。第二步配置训练参数TrainingArguments。这是框架的“中枢神经”所有行为都由此驱动training_args TrainingArguments( output_dir./bert-thucnews-checkpoints, # 模型和日志保存路径 num_train_epochs3, # 训练3轮 per_device_train_batch_size16, # 每张GPU的batch size per_device_eval_batch_size16, # 每张GPU的评估batch size warmup_ratio0.1, # 预热比例自动计算steps learning_rate2e-5, # 学习率BERT微调的经典值 weight_decay0.01, # 权重衰减防止过拟合 logging_strategysteps, # 日志策略按步记录 logging_steps100, # 每100步记录一次 evaluation_strategyepoch, # 评估策略每轮结束评估 save_strategyepoch, # 保存策略每轮结束保存 save_total_limit2, # 最多保存2个checkpoint load_best_model_at_endTrue, # 训练结束后加载最佳模型 metric_for_best_modeleval_f1, # 用f1分数判断最佳模型 greater_is_betterTrue, # f1越大越好 report_tonone, # 不上报到WB等平台本地日志 seed42, # 随机种子保证可复现 fp16True, # 启用混合精度训练速度翻倍 dataloader_num_workers4, # 数据加载器使用4个子进程 remove_unused_columnsFalse, # 不移除未使用的列因为我们只有input_ids等 )warmup_ratio0.1比warmup_steps500更智能它会根据总训练步数自动计算预热步数避免了手动计算的错误。fp16True是性能关键它让训练速度提升约40%且对精度影响极小。remove_unused_columnsFalse是安全选项虽然我们的数据集很干净但开启它会移除text等非模型输入列可能导致Trainer找不到labels而报错。第三步定义评估指标compute_metrics。Trainer需要一个函数来计算评估结果import numpy as np from sklearn.metrics import accuracy_score, f1_score, classification_report def compute_metrics(eval_pred): 计算评估指标 predictions, labels eval_pred # predictions是logits取argmax得到预测类别 preds np.argmax(predictions, axis1) # 计算accuracy和macro-f1 acc accuracy_score(labels, preds) f1 f1_score(labels, preds, averagemacro) return { accuracy: acc, f1: f1, eval_accuracy: acc, # Trainer会自动加eval_前缀 eval_f1: f1 }Trainer会将predict()的输出logits和tokenized_val中的labels传给这个函数。averagemacro是对所有类别f1的无偏平均比micro更能反映模型在少数类上的表现。至此模型、数据、参数、指标四大要素齐备只待Trainer登场。3.4 训练器Trainer初始化与训练启动全模块串联引擎Trainer是整个框架的指挥官它的初始化过程就是所有模块第一次正式“握手”的时刻。我们来逐行解析其初始化逻辑trainer Trainer( modelmodel, # 模型对象BertForSequenceClassification argstraining_args, # 训练参数TrainingArguments train_datasettokenized_train, # 训练数据集Dataset eval_datasettokenized_val, # 评估数据集Dataset tokenizertokenizer, # 分词器用于动态padding和生成 data_collatordata_collator, # 数据整理器DataCollatorWithPadding compute_metricscompute_metrics, # 评估函数callable )这短短几行背后是11步严格的初始化流程源码注释明确列出。第一步Args seedtraining_args被校验seed42被用来设置random、numpy、torch的随机种子确保结果可复现。第二步Accelerator loggingAccelerator被初始化它会自动检测设备CUDA/NPU/CPU并设置日志级别。第三步Model resolutionmodel被检查确认其forward()方法签名Liger Kernel等优化被标记。第四步Distributed strategy检查是否启用FSDP或DeepSpeed我们没配所以跳过。第五步Device placementmodel.to(args.device)被调用模型被移到GPU上。第六步Model introspectionTrainer分析模型发现它有labels输入和loss输出因此compute_loss()会直接取outputs[loss]。第七步Store init arguments所有传入的参数train_dataset,eval_dataset,compute_metrics被存储。第八步CallbacksDefaultFlowCallback、ProgressCallback等默认回调被注册。第九步Hub outputoutput_dir被创建如果启用了Hugging Face Hub会初始化repo。第十步Training stateTrainerState和TrainerControl被创建global_step0,epoch0。第十一步Finalizeuse_cacheTrue被设置XLA等特殊配置被处理。所有这些都在trainer Trainer(...)这一行内完成。初始化完成后启动训练# 开始训练 trainer.train() # 训练结束后保存最终模型 trainer.save_model(./bert-thucnews-final) # 也可以保存tokenizer方便后续推理 tokenizer.save_pretrained(./bert-thucnews-final)trainer.train()的执行会触发_inner_training_loop()它会准备阶段调用get_train_dataloader()DataCollator开始工作LengthGroupedSampler排序索引。训练阶段进入for epoch in range(num_train_epochs):循环对每个epoch调用_run_epoch()。单步执行在_run_epoch()中for update_step in range(num_update_steps_per_epoch):循环每次调用training_step()。损失计算training_step()调用compute_loss()后者调用model(**inputs)得到outputs[loss]。梯度更新loss.backward()optimizer.step()lr_scheduler.step()。回调触发on_step_end()被调用DefaultFlowCallback检查是否该log、evaluate、save。评估执行当control.should_evaluateTrueevaluate()被调用prediction_step()执行compute_metrics()计算分数。收尾工作_finalize_training()被调用清理资源保存最终状态。 整个过程model、data_collator、tokenizer、compute_metrics、callbacks全部被串联起来形成一个闭环。训练日志会显示类似Step 100/3000: loss0.4567, lr2e-05这行日志的背后是至少10个模块的协同作战。4. 关键技术点深度解析从LoRA微调到模块协作原理4.1 LoRA微调在BERT上实现参数高效迁移学习LoRALow-Rank Adaptation是当前大模型微调的主流范式它不修改原始BERT的庞大参数109M而是在其nn.Linear层旁“挂载”两个小矩阵A和B让梯度只在这两个小矩阵上更新。这使得微调显存占用从16GB骤降至4GB且效果媲美全参数微调。在Transformers中集成LoRA核心是peft库。第一步安装与导入pip install peft第二步配置LoRA参数。LoraConfig定义了挂载的位置和规模from peft import LoraConfig, get_peft_model lora_config LoraConfig( r8, # LoRA矩阵的秩8是常用值 lora_alpha16, # 缩放因子通常为2*r target_modules[query, value], # 在BERT的SelfAttention中只对query和value层挂载 lora_dropout0.1, # Dropout率防止过拟合 biasnone, # 不训练bias项 task_typeSEQ_CLS # 任务类型序列分类 )target_modules[query, value]是关键。BERT的BertSelfAttention层有query,key,value,dense四个nn.Linear子层。key层主要负责计算注意力权重dense层负责投影回隐藏空间它们对下游任务的贡献相对稳定因此我们只在query决定关注什么和value决定提取什么信息上做适配这是经验之选。第三步将LoRA应用到模型# 将LoRA配置应用到BERT模型上 peft_model get_peft_model(model, lora_config) peft_model.print_trainable_parameters() # 输出: trainable params: 1,048,576 || all params: 109,486,080 || trainable%: 0.9577get_peft_model()会遍历model的所有模块对匹配target_modules的nn.Linear层插入lora_A和lora_B两个可训练参数并重写其forward()方法original_output (x lora_A lora_B) * scaling。print_trainable_parameters()显示可训练参数仅约100万占总参数的0.96%这就是LoRA的威力。第四步与Trainer集成。peft_model是一个标准的nn.Module可以直接传给Trainertrainer Trainer( modelpeft_model, # 注意这里传的是peft_model不是原始model argstraining_args, train_datasettokenized_train, eval_datasettokenized_val, tokenizertokenizer, data_collatordata_collator, compute_metricscompute_metrics, ) trainer.train()训练完成后trainer.save_model()保存的是一个包含原始BERT权重和LoRA增量的完整模型。如果你想只保存LoRA权重体积更小便于分享可以peft_model.save_pretrained(./bert-thucnews-lora)这会生成adapter_config.json和adapter_model.bin两个文件。后续推理时只需加载原始BERT再用PeftModel.from_pretrained()加载LoRA权重即可。LoRA的“串联”体现在它无缝嵌入了Trainer的training_step()和compute_loss()流程中peft_model.forward()被调用其内部的lora_A和lora_B参与计算梯度只反向传播到这两个小矩阵上optimizer也只更新它们。整个过程对Trainer完全透明这就是Transformers框架“约定优于配置”的绝佳体现。4.2 模块协作原理Trainer、Callback与Optimizer的三角关系Trainer、Callback和Optimizer构成了Transformers训练的“铁三角”它们的协作不是简单的调用关系而是一种基于事件驱动的松耦合。理解这个三角是掌握框架灵魂的关键。Trainer是中心节点它定义了训练的主干流程train()→_inner_training_loop()→_run_epoch()→training_step()。Callback是观察者它不主动做事而是被动等待Trainer在特定时刻发出的“事件信号”。Optimizer是执行者它只关心如何根据梯度更新参数对训练流程一无所知。它们的协作通过TrainerControl这个“信号灯”来实现。我们以早停Early Stopping为例看这个三角如何联动。首先EarlyStoppingCallback被注册到Trainer中from transformers import EarlyStoppingCallback trainer Trainer( # ... 其他参数 callbacks[EarlyStoppingCallback(early_stopping_patience3)] )Trainer在初始化时会将这个回调加入callback_handler.callbacks列表。当训练进行到某个step_inner_training_loop()执行完training_step()后会调用callback_handler.on_step_end()。这个方法会遍历所有回调对每个回调调用其on_step_end()方法。EarlyStoppingCallback.on_step_end()什么都不做因为它只关心评估结果。当control.should_evaluateTrue时Trainer调用evaluate()并在评估完成后调用callback_handler.on_evaluate()。此时EarlyStoppingCallback.on_evaluate()被触发def on_evaluate(self, args, state, control, metrics, **kwargs): # metrics是compute_metrics返回的字典如{eval_f1: 0.85} metric_value metrics.get(self.metric_for_best_model, 0) # 检查是否比历史最佳值更好 if self.best_metric is None or self.is_better(metric_value, self.best_metric): self.best_metric metric_value self.early_stopping_patience_counter 0 else: self.early_stopping_patience_counter 1 # 如果连续patience次未提升则发送停止信号 if self.early_stopping_patience_counter self.early_stopping_patience: control.should_training_stop True # 关键修改control对象 return control # 必须返回controlTrainer会接收注意最后一行return control。callback_handler.on_evaluate()会接收这个被修改过的control对象并将其返回给Trainer。Trainer拿到这个control后检查control.should_training_stop如果为True则立即跳出for epoch循环终止训练。整个过程Trainer没有调用EarlyStoppingCallback的任何私有方法EarlyStoppingCallback也没有调用Trainer的任何私有属性

相关新闻