【Bug已解决】Save only best weights with huggingface transformers 解决方案

问题描述

在使用 Hugging Face Transformers 库进行模型微调（fine-tuning）时，一个常见的需求是只保存验证集上表现最好的模型权重，而不是保存每个 epoch 结束后的权重。这种"最佳模型保存"策略可以：

• 避免保存大量中间模型浪费磁盘空间
• 确保最终使用的是性能最优的模型
• 防止过拟合后使用性能下降的模型

然而，许多开发者在使用TrainerAPI 或自定义训练循环时，不知道如何正确实现"只保存最佳权重"的功能，遇到以下问题：

1. Trainer默认每个 epoch 都保存模型，磁盘空间迅速耗尽
2. 不知道如何设置load_best_model_at_end参数
3. 自定义训练循环中手动保存最佳模型的逻辑有 bug
4. EarlyStoppingCallback与最佳模型保存的配合问题
5. 保存的"最佳模型"实际上不是最佳的（评估指标方向错误）
6. 多个指标如何选择最佳模型

错误复现

场景一：Trainer 保存了所有 epoch 的模型

from transformers import Trainer, TrainingArguments, AutoModelForSequenceClassification model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2) training_args = TrainingArguments( output_dir='./results', num_train_epochs=10, save_strategy='epoch', # 每个 epoch 都保存 ) trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, ) trainer.train() # 结果：results/ 目录下有 checkpoint-xxx 多个子目录，每个都很大

场景二：load_best_model_at_end配置错误

training_args = TrainingArguments( output_dir='./results', num_train_epochs=10, save_strategy='epoch', load_best_model_at_end=True, # 启用最佳模型加载 # 但没有设置 metric_for_best_model ) # 报错或行为不符合预期

场景三：自定义训练循环中最佳模型保存有 bug

best_score = 0 for epoch in range(num_epochs): # 训练... # 评估 eval_score = evaluate(model, eval_dataloader) if eval_score > best_score: best_score = eval_score torch.save(model.state_dict(), 'best_model.pth') # 问题：如果指标是 loss（越小越好），这个逻辑是错误的

场景四：EarlyStopping 与最佳模型保存冲突

from transformers import EarlyStoppingCallback training_args = TrainingArguments( output_dir='./results', num_train_epochs=100, save_strategy='epoch', load_best_model_at_end=True, metric_for_best_model='eval_loss', greater_is_better=False, # loss 越小越好 ) trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, callbacks=[EarlyStoppingCallback(early_stopping_patience=3)], ) trainer.train() # 问题：EarlyStopping 停止后，加载的模型可能不是最佳的

根因分析

1. Trainer 的保存策略

Hugging FaceTrainer默认的保存策略由save_strategy控制：

• 'no'：不保存
• 'epoch'：每个 epoch 结束保存
• 'steps'：每 N 步保存

默认情况下，Trainer会保存完整的 checkpoint（模型权重、优化器状态、调度器状态等），每个 checkpoint 可能几百 MB 到几 GB。

2.load_best_model_at_end的工作机制

当load_best_model_at_end=True时，Trainer会在训练结束时加载最佳模型。但这需要：

• save_strategy与evaluation_strategy一致
• 设置metric_for_best_model指定评估指标
• 设置greater_is_better指定指标方向

3. 指标方向的重要性

不同的评估指标有不同的"好"的方向：

• accuracy、F1、precision、recall：越大越好（greater_is_better=True）
• loss、perplexity、WER：越小越好（greater_is_better=False）

如果方向设置错误，保存的"最佳模型"实际上是最差的。

4. Checkpoint 清理

Trainer默认会保存所有 checkpoint（save_total_limit=None）。需要设置save_total_limit=1来只保留最新的（或最佳的）checkpoint。

5. 自定义训练循环的复杂性

在自定义训练循环中实现最佳模型保存需要考虑：

• 指标方向（越大越好还是越小越好）
• 保存内容（state_dict vs 完整模型）
• 恢复训练（是否保存优化器状态）
• 多指标选择

解决方案

方案一：使用 Trainer 的load_best_model_at_end

from transformers import ( Trainer, TrainingArguments, AutoModelForSequenceClassification, EarlyStoppingCallback ) model = AutoModelForSequenceClassification.from_pretrained( 'bert-base-uncased', num_labels=2 ) training_args = TrainingArguments( output_dir='./results', num_train_epochs=10, # 保存和评估策略必须一致 save_strategy='epoch', evaluation_strategy='epoch', # 加载最佳模型 load_best_model_at_end=True, metric_for_best_model='eval_loss', greater_is_better=False, # loss 越小越好 # 只保留 1 个 checkpoint save_total_limit=1, # 输出日志 logging_dir='./logs', logging_steps=10, ) trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, compute_metrics=compute_metrics, # 自定义指标函数 callbacks=[EarlyStoppingCallback(early_stopping_patience=3)], ) trainer.train() # 训练结束后，model 已经是最佳模型 # 保存最佳模型 trainer.save_model('./best_model')

方案二：自定义训练循环中保存最佳模型

import torch import os def train_with_best_model_saving( model, train_dataloader, eval_dataloader, optimizer, num_epochs, device, metric_name='loss', greater_is_better=False, save_dir='./best_model' ): """自定义训练循环，只保存最佳模型""" if greater_is_better: best_score = float('-inf') compare = lambda current, best: current > best else: best_score = float('inf') compare = lambda current, best: current < best os.makedirs(save_dir, exist_ok=True) for epoch in range(num_epochs): # 训练 model.train() train_loss = 0 for batch in train_dataloader: batch = {k: v.to(device) for k, v in batch.items()} optimizer.zero_grad() outputs = model(**batch) loss = outputs.loss loss.backward() optimizer.step() train_loss += loss.item() # 评估 model.eval() eval_loss = 0 eval_score = 0 with torch.no_grad(): for batch in eval_dataloader: batch = {k: v.to(device) for k, v in batch.items()} outputs = model(**batch) eval_loss += outputs.loss.item() # 计算自定义指标... eval_loss /= len(eval_dataloader) current_score = eval_loss if metric_name == 'loss' else eval_score # 判断是否是最佳模型 if compare(current_score, best_score): best_score = current_score # 保存最佳模型 torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'best_score': best_score, }, os.path.join(save_dir, 'best_model.pth')) print(f"Epoch {epoch}: New best {metric_name} = {best_score:.4f}, model saved!") else: print(f"Epoch {epoch}: {metric_name} = {current_score:.4f} (best: {best_score:.4f})") # 加载最佳模型 checkpoint = torch.load(os.path.join(save_dir, 'best_model.pth')) model.load_state_dict(checkpoint['model_state_dict']) print(f"\nLoaded best model from epoch {checkpoint['epoch']} with {metric_name} = {checkpoint['best_score']:.4f}") return model

方案三：使用多指标选择最佳模型

from transformers import TrainerCallback class BestModelTracker(TrainerCallback): """跟踪最佳模型的自定义回调""" def __init__(self, metric_name='eval_loss', greater_is_better=False): self.metric_name = metric_name self.greater_is_better = greater_is_better self.best_score = float('-inf') if greater_is_better else float('inf') self.best_epoch = 0 def on_evaluate(self, args, state, control, metrics=None, **kwargs): if metrics is None or self.metric_name not in metrics: return current_score = metrics[self.metric_name] is_better = ( current_score > self.best_score if self.greater_is_better else current_score < self.best_score ) if is_better: self.best_score = current_score self.best_epoch = state.epoch print(f"New best {self.metric_name}: {self.best_score:.4f} at epoch {self.best_epoch:.1f}")

完整修复代码

以下是一个完整的、生产可用的最佳模型保存方案：

""" Hugging Face Transformers 最佳模型保存完整方案 支持 Trainer API 和自定义训练循环 """ import torch import os import numpy as np from typing import Dict, Optional, Callable from dataclasses import dataclass # Transformers 相关导入 try: from transformers import ( Trainer, TrainingArguments, AutoModelForSequenceClassification, AutoTokenizer, EarlyStoppingCallback, TrainerCallback, ) from transformers.trainer_utils import EvalPrediction TRANSFORMERS_AVAILABLE = True except ImportError: TRANSFORMERS_AVAILABLE = False print("Warning: transformers not installed. Some features unavailable.") # ============ 方案1: Trainer API ============ def train_with_trainer_best_model( model_name: str = 'bert-base-uncased', num_labels: int = 2, output_dir: str = './results', best_model_dir: str = './best_model', num_train_epochs: int = 10, learning_rate: float = 2e-5, per_device_train_batch_size: int = 16, per_device_eval_batch_size: int = 32, metric_for_best_model: str = 'eval_loss', greater_is_better: bool = False, early_stopping_patience: int = 3, train_dataset=None, eval_dataset=None, compute_metrics: Optional[Callable] = None, ): """ 使用 Trainer API 训练并保存最佳模型 Args: model_name: 预训练模型名称 num_labels: 分类标签数 output_dir: 训练输出目录 best_model_dir: 最佳模型保存目录 num_train_epochs: 训练轮数 metric_for_best_model: 用于选择最佳模型的指标 greater_is_better: 指标越大越好时为 True early_stopping_patience: 早停耐心值 """ if not TRANSFORMERS_AVAILABLE: raise ImportError("transformers 库未安装") print("=" * 60) print("使用 Trainer API 训练并保存最佳模型") print("=" * 60) # 加载模型 model = AutoModelForSequenceClassification.from_pretrained( model_name, num_labels=num_labels ) # 配置训练参数 training_args = TrainingArguments( output_dir=output_dir, # 训练参数 num_train_epochs=num_train_epochs, learning_rate=learning_rate, per_device_train_batch_size=per_device_train_batch_size, per_device_eval_batch_size=per_device_eval_batch_size, weight_decay=0.01, warmup_ratio=0.1, # 保存和评估策略（必须一致） save_strategy='epoch', evaluation_strategy='epoch', # 最佳模型配置 load_best_model_at_end=True, metric_for_best_model=metric_for_best_model, greater_is_better=greater_is_better, # Checkpoint 清理 save_total_limit=1, # 只保留 1 个 checkpoint # 日志 logging_dir=f'{output_dir}/logs', logging_steps=50, report_to='tensorboard', # 其他 seed=42, dataloader_num_workers=4, ) # 创建 Trainer trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, compute_metrics=compute_metrics, callbacks=[ EarlyStoppingCallback( early_stopping_patience=early_stopping_patience, early_stopping_threshold=0.001 ) ] if early_stopping_patience > 0 else [], ) # 训练 trainer.train() # 训练结束后，model 已是最佳模型（load_best_model_at_end=True） # 保存最佳模型 trainer.save_model(best_model_dir) print(f"\n最佳模型已保存到 {best_model_dir}") # 评估最佳模型 eval_results = trainer.evaluate() print(f"\n最佳模型评估结果:") for key, value in eval_results.items(): print(f" {key}: {value:.4f}") return trainer, model # ============ 方案2: 自定义训练循环 ============ @dataclass class BestModelConfig: """最佳模型保存配置""" metric_name: str = 'loss' greater_is_better: bool = False save_dir: str = './best_model' save_total_limit: int = 1 # 保留几个最佳模型 verbose: bool = True class BestModelSaver: """最佳模型保存器""" def __init__(self, config: BestModelConfig): self.config = config self.best_score = float('-inf') if config.greater_is_better else float('inf') self.best_epoch = -1 self.saved_models = [] # 保存的模型文件列表 os.makedirs(config.save_dir, exist_ok=True) def _is_better(self, current: float) -> bool: """判断当前分数是否更好""" if self.config.greater_is_better: return current > self.best_score else: return current < self.best_score def save_if_best( self, model: torch.nn.Module, current_score: float, epoch: int, optimizer: Optional[torch.optim.Optimizer] = None, extra_info: Optional[Dict] = None, ) -> bool: """ 如果是最佳模型则保存 Returns: 是否保存了新模型 """ if not self._is_better(current_score): if self.config.verbose: print(f" Epoch {epoch}: {self.config.metric_name}={current_score:.4f} " f"(best: {self.best_score:.4f} at epoch {self.best_epoch})") return False # 更新最佳分数 old_best = self.best_score self.best_score = current_score self.best_epoch = epoch # 保存模型 filepath = os.path.join( self.config.save_dir, f'best_model_epoch{epoch}_{self.config.metric_name}{current_score:.4f}.pth' ) checkpoint = { 'epoch': epoch, 'model_state_dict': model.state_dict(), 'best_score': current_score, 'metric_name': self.config.metric_name, } if optimizer is not None: checkpoint['optimizer_state_dict'] = optimizer.state_dict() if extra_info: checkpoint.update(extra_info) torch.save(checkpoint, filepath) self.saved_models.append(filepath) # 清理多余的模型 while len(self.saved_models) > self.config.save_total_limit: old_file = self.saved_models.pop(0) if os.path.exists(old_file): os.remove(old_file) if self.config.verbose: print(f" Removed old checkpoint: {os.path.basename(old_file)}") if self.config.verbose: print(f" Epoch {epoch}: New best {self.config.metric_name}={current_score:.4f} " f"(was: {old_best:.4f}), saved to {os.path.basename(filepath)}") return True def load_best_model(self, model: torch.nn.Module, device: str = 'cpu') -> Dict: """加载最佳模型""" if not self.saved_models: raise ValueError("没有保存的模型") filepath = self.saved_models[-1] # 最新的（最佳的） checkpoint = torch.load(filepath, map_location=device) model.load_state_dict(checkpoint['model_state_dict']) if self.config.verbose: print(f"Loaded best model from epoch {checkpoint['epoch']} " f"with {self.config.metric_name}={checkpoint['best_score']:.4f}") return checkpoint def get_best_score(self) -> float: """获取最佳分数""" return self.best_score def get_best_epoch(self) -> int: """获取最佳 epoch""" return self.best_epoch def train_with_custom_loop( model: torch.nn.Module, train_dataloader: torch.utils.data.DataLoader, eval_dataloader: torch.utils.data.DataLoader, optimizer: torch.optim.Optimizer, num_epochs: int, device: torch.device, config: BestModelConfig, eval_fn: Optional[Callable] = None, ): """ 自定义训练循环，保存最佳模型 Args: model: PyTorch 模型 train_dataloader: 训练数据加载器 eval_dataloader: 评估数据加载器 optimizer: 优化器 num_epochs: 训练轮数 device: 计算设备 config: 最佳模型保存配置 eval_fn: 自定义评估函数（返回指标字典） """ print("=" * 60) print("自定义训练循环 - 最佳模型保存") print("=" * 60) print(f"指标: {config.metric_name}, 越大越好: {config.greater_is_better}") print(f"保存目录: {config.save_dir}") print() saver = BestModelSaver(config) for epoch in range(num_epochs): # ========== 训练阶段 ========== model.train() total_loss = 0 num_batches = 0 for batch in train_dataloader: if isinstance(batch, dict): batch = {k: v.to(device) for k, v in batch.items()} else: batch = batch.to(device) optimizer.zero_grad() if isinstance(batch, dict): outputs = model(**batch) else: outputs = model(batch) loss = outputs.loss if hasattr(outputs, 'loss') else outputs[0] loss.backward() optimizer.step() total_loss += loss.item() num_batches += 1 train_loss = total_loss / num_batches # ========== 评估阶段 ========== model.eval() if eval_fn is not None: # 使用自定义评估函数 metrics = eval_fn(model, eval_dataloader, device) else: # 默认使用 loss 作为指标 eval_loss = 0 eval_batches = 0 with torch.no_grad(): for batch in eval_dataloader: if isinstance(batch, dict): batch = {k: v.to(device) for k, v in batch.items()} else: batch = batch.to(device) if isinstance(batch, dict): outputs = model(**batch) else: outputs = model(batch) loss = outputs.loss if hasattr(outputs, 'loss') else outputs[0] eval_loss += loss.item() eval_batches += 1 metrics = {'loss': eval_loss / eval_batches} # ========== 保存最佳模型 ========== current_score = metrics.get(config.metric_name, metrics.get(f'eval_{config.metric_name}', 0)) print(f"Epoch {epoch+1}/{num_epochs}: train_loss={train_loss:.4f}", end="") for k, v in metrics.items(): print(f", {k}={v:.4f}", end="") print() saver.save_if_best( model=model, current_score=current_score, epoch=epoch, optimizer=optimizer, extra_info={'train_loss': train_loss, 'all_metrics': metrics} ) # ========== 加载最佳模型 ========== print("\n" + "=" * 60) print("训练完成，加载最佳模型") print("=" * 60) checkpoint = saver.load_best_model(model, device=str(device)) print(f"最佳 epoch: {checkpoint['epoch']}") print(f"最佳 {config.metric_name}: {checkpoint['best_score']:.4f}") print(f"训练 loss: {checkpoint.get('train_loss', 'N/A')}") if 'all_metrics' in checkpoint: print("所有指标:") for k, v in checkpoint['all_metrics'].items(): print(f" {k}: {v:.4f}") return model, checkpoint # ============ 使用示例 ============ def demo_custom_loop(): """演示自定义训练循环""" print("\n" + "=" * 60) print("演示: 自定义训练循环") print("=" * 60) # 创建简单模型和数据 import torch.nn as nn from torch.utils.data import DataLoader, TensorDataset model = nn.Sequential( nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10) ) # 模拟数据 X_train = torch.randn(1000, 784) y_train = torch.randint(0, 10, (1000,)) X_eval = torch.randn(200, 784) y_eval = torch.randint(0, 10, (200,)) train_dataset = TensorDataset(X_train, y_train) eval_dataset = TensorDataset(X_eval, y_eval) train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True) eval_loader = DataLoader(eval_dataset, batch_size=32) optimizer = torch.optim.Adam(model.parameters(), lr=0.001) device = torch.device('cpu') # 自定义评估函数 def eval_fn(model, dataloader, device): model.eval() correct = 0 total = 0 loss = 0 criterion = nn.CrossEntropyLoss() with torch.no_grad(): for X, y in dataloader: X, y = X.to(device), y.to(device) outputs = model(X) loss += criterion(outputs, y).item() _, predicted = outputs.max(1) correct += predicted.eq(y).sum().item() total += y.size(0) return { 'loss': loss / len(dataloader), 'accuracy': correct / total } # 配置 config = BestModelConfig( metric_name='accuracy', greater_is_better=True, # accuracy 越大越好 save_dir='./demo_best_model', save_total_limit=1, verbose=True, ) # 训练 model, checkpoint = train_with_custom_loop( model=model, train_dataloader=train_loader, eval_dataloader=eval_loader, optimizer=optimizer, num_epochs=10, device=device, config=config, eval_fn=eval_fn, ) # 清理 import shutil if os.path.exists('./demo_best_model'): shutil.rmtree('./demo_best_model') def main(): print("\n" + "=" * 60) print(" Hugging Face Transformers 最佳模型保存方案") print("=" * 60) demo_custom_loop() print("\n" + "=" * 60) print(" 演示完成！") print("=" * 60) if __name__ == '__main__': main()

常见陷阱与注意事项

1.save_strategy和evaluation_strategy必须一致

当使用load_best_model_at_end=True时，save_strategy和evaluation_strategy必须一致（都是'epoch'或都是'steps'）。否则Trainer无法在评估时找到对应的 checkpoint。

# 正确 TrainingArguments(save_strategy='epoch', evaluation_strategy='epoch', ...) # 错误 TrainingArguments(save_strategy='steps', evaluation_strategy='epoch', ...) # 报错: save_strategy and evaluation_strategy must match

2.metric_for_best_model的命名

metric_for_best_model的值需要与评估指标名称匹配。如果使用compute_metrics函数返回的指标，名称不需要加eval_前缀（Trainer 会自动添加）：

# compute_metrics 返回 {'accuracy': 0.95, 'f1': 0.93} # metric_for_best_model 应该是 'accuracy' 或 'f1'（不带 eval_ 前缀） TrainingArguments(metric_for_best_model='accuracy', ...)

3.greater_is_better的方向

常见的指标方向：

• 越大越好（greater_is_better=True）：accuracy, f1, precision, recall, auc
• 越小越好（greater_is_better=False）：loss, perplexity, wer, cer

如果不确定，可以查看指标的文档或使用常识判断。

4.save_total_limit与load_best_model_at_end的交互

当save_total_limit=1且load_best_model_at_end=True时，Trainer会确保保留最佳模型的 checkpoint。但如果save_total_limit设置为 0，可能会删除所有 checkpoint 导致无法加载最佳模型。

5. EarlyStopping 后的模型加载

EarlyStoppingCallback会在指标不再改善时停止训练。配合load_best_model_at_end=True，训练结束后会加载停止前的最佳模型（而不是最后一个 epoch 的模型）。

6. 自定义训练循环中的指标方向

在自定义训练循环中，必须正确处理指标方向。使用BestModelSaver类可以避免手动比较的错误。

7. 保存完整 checkpoint vs 仅 state_dict

• 仅 state_dict：文件小，加载时需要模型定义代码
• 完整 checkpoint（含优化器状态）：文件大，可以恢复训练

对于只需要推理的场景，保存 state_dict 即可。对于可能需要恢复训练的场景，保存完整 checkpoint。

8. 多 GPU 训练中的模型保存

在 DataParallel 或 DistributedDataParallel 下，保存模型时需要保存model.module.state_dict()而不是model.state_dict()，以避免键名前缀问题。

9.Trainer的save_model方法

trainer.save_model()保存的是当前内存中的模型状态。如果load_best_model_at_end=True，训练结束后模型已经是最佳状态，直接调用save_model()即可保存最佳模型。

10. 验证保存的模型

保存最佳模型后，建议重新加载并评估，确保保存的模型确实是最佳的：

# 重新加载并评估 from transformers import AutoModelForSequenceClassification loaded_model = AutoModelForSequenceClassification.from_pretrained('./best_model') # 评估 loaded_model，确认指标与训练时记录的最佳值一致

总结

在 Hugging Face Transformers 中实现"只保存最佳模型权重"的功能，有两种主要方式：

方式一：使用 Trainer API（推荐）

TrainingArguments( load_best_model_at_end=True, metric_for_best_model='eval_loss', greater_is_better=False, save_total_limit=1, ) trainer.train() trainer.save_model('./best_model')

方式二：自定义训练循环使用BestModelSaver类，在每次评估后检查是否是最佳模型，是则保存。

关键要点：

1. 正确设置指标方向（greater_is_better）
2. 确保保存和评估策略一致
3. 使用save_total_limit限制 checkpoint 数量
4. 配合EarlyStoppingCallback防止过拟合
5. 训练结束后验证最佳模型的性能

通过本文提供的完整代码，可以轻松实现最佳模型保存功能，无论是使用 Trainer API 还是自定义训练循环。