【Bug已解决】Saving and reload huggingface fine-tuned transformer 解决方案问题描述在使用 Hugging Face Transformers 库对预训练模型进行微调fine-tuning后许多开发者会遇到一个常见且令人头疼的问题如何正确地保存微调后的模型并在后续使用时完整地重新加载它。这个看似简单的操作实际上涉及多个层面的复杂性包括模型权重、tokenizer、配置文件、训练状态等的持久化。典型的报错场景包括保存后重新加载时模型输出全是NaN或预测结果完全随机加载时报错AttributeError: XXX object has no attribute XXXTokenizer 与模型不匹配导致输入处理不一致使用torch.save()保存后跨环境加载出现ModuleNotFoundError保存的模型文件过大或过小怀疑保存不完整这些问题的根本原因在于 Hugging Face 模型的保存机制与普通 PyTorch 模型有所不同需要同时保存多个组件才能确保模型的完整性和可复现性。错误复现场景一仅保存模型权重导致加载失败# 错误做法只保存了 state_dict import torch from transformers import AutoModelForSequenceClassification, AutoTokenizer from torch.utils.data import DataLoader # 加载预训练模型 model AutoModelForSequenceClassification.from_pretrained( bert-base-uncased, num_labels2 ) tokenizer AutoTokenizer.from_pretrained(bert-base-uncased) # 模拟微调过程简化版 optimizer torch.optim.AdamW(model.parameters(), lr2e-5) model.train() for epoch in range(3): for batch in dataloader: outputs model(**batch) loss outputs.loss loss.backward() optimizer.step() optimizer.zero_grad() # 错误仅保存 state_dict torch.save(model.state_dict(), ./fine_tuned_model.pt) # 后续加载时出现问题 model_loaded AutoModelForSequenceClassification.from_pretrained( bert-base-uncased, num_labels2 ) model_loaded.load_state_dict(torch.load(./fine_tuned_model.pt)) # 问题tokenizer 没有保存配置可能不匹配 # 在新环境中可能报错或结果异常场景二跨环境加载报错# 在环境A中保存 torch.save(model, ./full_model.pt) # 在环境B中加载 model torch.load(./full_model.pt) # 报错ModuleNotFoundError: No module named transformers.models.bert.modeling_bert # 或AttributeError: Cannot get attribute XXX on module transformers...场景三保存后预测结果异常# 保存模型 model.save_pretrained(./my_model) tokenizer.save_pretrained(./my_model) # 加载模型 model AutoModelForSequenceClassification.from_pretrained(./my_model) # 但忘记加载 tokenizer使用了默认 tokenizer tokenizer AutoTokenizer.from_pretrained(bert-base-uncased) # 结果预测结果完全错误因为 tokenizer 配置可能不一致根因分析1. Hugging Face 模型的组件结构Hugging Face 的 Transformer 模型不是一个简单的nn.Module它由多个组件构成模型权重model weightspytorch_model.bin或model.safetensors模型配置config.json包含模型架构、超参数等信息Tokenizer 文件vocab.txt、tokenizer_config.json、special_tokens_map.json等生成配置generation_config.json对于生成模型训练状态optimizer state、scheduler state、random state 等仅保存state_dict会丢失配置信息导致加载时使用默认配置可能与微调时的配置不一致。2.torch.save()的序列化问题torch.save()使用 Python 的pickle模块进行序列化。当直接保存整个模型对象时# 这种方式保存的是模型的类路径和实例属性 torch.save(model, model.pt) # pickle 会记录: transformers.models.bert.modeling_bert.BertForSequenceClassification # 在新环境中如果版本不同类路径可能变化导致反序列化失败3. Tokenizer 的重要性Tokenizer 决定了文本如何被切分和编码。微调过程中如果修改了 tokenizer如添加特殊 token不保存 tokenizer 会导致加载后使用原始 tokenizer输入编码不一致。4. 模型评估模式未设置加载模型后忘记调用model.eval()导致 dropout 等层仍处于训练模式预测结果不确定。解决方案方案一使用save_pretrained()和from_pretrained()推荐这是 Hugging Face 官方推荐的方式会保存所有必要组件import torch from transformers import ( AutoModelForSequenceClassification, AutoTokenizer, TrainingArguments, Trainer, AutoConfig ) from datasets import load_dataset import numpy as np import os # # 完整的微调 保存 加载流程 # def fine_tune_and_save(): 完整的微调、保存和加载流程 model_name bert-base-uncased save_directory ./fine_tuned_bert # 1. 加载 tokenizer 和模型 tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForSequenceClassification.from_pretrained( model_name, num_labels2, # 自定义配置 hidden_dropout_prob0.1, attention_probs_dropout_prob0.1, ) print(f模型配置: {model.config}) print(f模型参数量: {sum(p.numel() for p in model.parameters()):,}) # 2. 准备数据集 dataset load_dataset(imdb, split{train: train[:1000], test: test[:200]}) def tokenize_function(examples): return tokenizer( examples[text], paddingmax_length, truncationTrue, max_length128, ) tokenized_datasets dataset.map(tokenize_function, batchedTrue) tokenized_datasets tokenized_datasets.remove_columns([text]) tokenized_datasets tokenized_datasets.rename_column(label, labels) tokenized_datasets.set_format(torch) # 3. 训练配置 training_args TrainingArguments( output_dir./training_output, num_train_epochs3, per_device_train_batch_size8, per_device_eval_batch_size16, warmup_steps100, weight_decay0.01, logging_dir./logs, logging_steps50, evaluation_strategyepoch, save_strategyepoch, load_best_model_at_endTrue, metric_for_best_modelaccuracy, ) # 4. 定义评估指标 def compute_metrics(eval_pred): logits, labels eval_pred predictions np.argmax(logits, axis-1) accuracy (predictions labels).mean() return {accuracy: accuracy} # 5. 创建 Trainer 并训练 trainer Trainer( modelmodel, argstraining_args, train_datasettokenized_datasets[train], eval_datasettokenized_datasets[test], compute_metricscompute_metrics, ) trainer.train() # 6. 评估 eval_results trainer.evaluate() print(f评估结果: {eval_results}) # # 关键正确保存模型和 tokenizer # # 创建保存目录 os.makedirs(save_directory, exist_okTrue) # 保存模型包含权重和配置 model.save_pretrained(save_directory) # 保存 tokenizer包含词表和配置 tokenizer.save_pretrained(save_directory) # 验证保存的文件 saved_files os.listdir(save_directory) print(f保存的文件: {saved_files}) # 预期文件列表: # - config.json (模型配置) # - model.safetensors 或 pytorch_model.bin (模型权重) # - vocab.txt (词表) # - tokenizer_config.json (tokenizer配置) # - special_tokens_map.json (特殊token映射) # - tokenizer.json (fast tokenizer文件) return save_directory def load_and_predict(save_directory): 正确加载模型并进行预测 # 1. 加载 tokenizer必须从同一目录加载 tokenizer AutoTokenizer.from_pretrained(save_directory) # 2. 加载模型自动读取 config.json model AutoModelForSequenceClassification.from_pretrained(save_directory) # 3. 设置为评估模式关键 model.eval() # 4. 如果有 GPU移到 GPU device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) # 5. 进行预测 test_texts [ This movie is absolutely fantastic! Highly recommended., Terrible film. Complete waste of time and money., An average movie, nothing special but not bad either., ] # 编码输入 inputs tokenizer( test_texts, paddingTrue, truncationTrue, max_length128, return_tensorspt ).to(device) # 推理 with torch.no_grad(): outputs model(**inputs) predictions torch.nn.functional.softmax(outputs.logits, dim-1) predicted_labels torch.argmax(predictions, dim-1) # 输出结果 labels [负面, 正面] for text, pred, probs in zip(test_texts, predicted_labels, predictions): print(f文本: {text[:50]}...) print(f预测: {labels[pred.item()]} (置信度: {probs[pred].item():.4f})) print() # 执行完整流程 if __name__ __main__: save_dir fine_tune_and_save() load_and_predict(save_dir)方案二保存完整训练状态用于断点续训如果需要中断后继续训练需要保存更多状态import torch from transformers import AutoModelForSequenceClassification, AutoTokenizer import os import json from datetime import datetime def save_training_checkpoint(model, tokenizer, optimizer, scheduler, epoch, step, save_dir): 保存完整的训练检查点支持断点续训 os.makedirs(save_dir, exist_okTrue) # 1. 保存模型和 tokenizer使用 Hugging Face 方法 model.save_pretrained(save_dir) tokenizer.save_pretrained(save_dir) # 2. 保存训练状态使用 PyTorch 方法 checkpoint { epoch: epoch, step: step, optimizer_state_dict: optimizer.state_dict(), scheduler_state_dict: scheduler.state_dict() if scheduler else None, random_state: { python: random.getstate(), numpy: np.random.get_state(), torch: torch.get_rng_state(), cuda: torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None, }, timestamp: datetime.now().isoformat(), } torch.save(checkpoint, os.path.join(save_dir, training_state.pt)) # 3. 保存训练元信息 meta_info { epoch: epoch, step: step, model_type: model.__class__.__name__, save_time: datetime.now().isoformat(), } with open(os.path.join(save_dir, training_meta.json), w) as f: json.dump(meta_info, f, indent2) print(f训练检查点已保存到 {save_dir}) print(f - Epoch: {epoch}, Step: {step}) def load_training_checkpoint(save_dir, model_nameNone): 加载训练检查点恢复训练状态 # 1. 加载模型和 tokenizer model AutoModelForSequenceClassification.from_pretrained(save_dir) tokenizer AutoTokenizer.from_pretrained(save_dir) # 2. 加载训练状态 checkpoint torch.load(os.path.join(save_dir, training_state.pt)) # 3. 恢复随机状态确保可复现性 random.setstate(checkpoint[random_state][python]) np.random.set_state(checkpoint[random_state][numpy]) torch.set_rng_state(checkpoint[random_state][torch]) if checkpoint[random_state][cuda] is not None: torch.cuda.set_rng_state_all(checkpoint[random_state][cuda]) # 4. 创建 optimizer 并恢复状态 optimizer torch.optim.AdamW(model.parameters(), lr2e-5) optimizer.load_state_dict(checkpoint[optimizer_state_dict]) # 5. 恢复 scheduler如果有 scheduler None if checkpoint[scheduler_state_dict] is not None: from transformers import get_linear_schedule_with_warmup scheduler get_linear_schedule_with_warmup( optimizer, num_warmup_steps100, num_training_steps10000 ) scheduler.load_state_dict(checkpoint[scheduler_state_dict]) print(f训练检查点已从 {save_dir} 加载) print(f - 恢复到 Epoch: {checkpoint[epoch]}, Step: {checkpoint[step]})  return model, tokenizer, optimizer, scheduler, checkpoint[epoch], checkpoint[step]方案三使用 Safetensors 格式更安全from transformers import AutoModelForSequenceClassification # 保存为 safetensors 格式默认行为更安全 model AutoModelForSequenceClassification.from_pretrained(bert-base-uncased) # safetensors 格式避免了 pickle 的安全风险 # 且加载速度更快 model.save_pretrained( ./safe_model, safe_serializationTrue, # 使用 safetensors 格式 max_shard_size500MB, # 大模型分片保存 ) # 加载时自动识别格式 model AutoModelForSequenceClassification.from_pretrained(./safe_model)完整修复代码下面是一个完整的、可运行的端到端示例涵盖微调、保存、加载和验证 完整的 Hugging Face Transformer 微调模型保存与加载示例 包含数据准备、微调训练、模型保存、模型加载、结果验证 import torch import torch.nn as nn import numpy as np import random import os import json from datetime import datetime from typing import Dict, List, Tuple, Optional # Hugging Face 相关导入 from transformers import ( AutoModelForSequenceClassification, AutoTokenizer, AutoConfig, TrainingArguments, Trainer, get_linear_schedule_with_warmup, set_seed, ) from datasets import load_dataset, Dataset from torch.utils.data import DataLoader # # 配置类 # class FineTuneConfig: 微调配置 model_name: str bert-base-uncased num_labels: int 2 max_length: int 128 batch_size: int 8 eval_batch_size: int 16 learning_rate: float 2e-5 num_epochs: int 3 warmup_ratio: float 0.1 weight_decay: float 0.01 seed: int 42 save_dir: str ./fine_tuned_model_complete classmethod def to_dict(cls): return {k: v for k, v in cls.__dict__.items() if not k.startswith(_)} # # 数据准备模块 # class DataPreparation: 数据准备和预处理 def __init__(self, config: FineTuneConfig): self.config config self.tokenizer AutoTokenizer.from_pretrained(config.model_name) def prepare_data(self) - Tuple[Dataset, Dataset]: 准备训练和评估数据集 # 使用 IMDB 数据集的子集 raw_datasets load_dataset(imdb, split{ train: ftrain[:{1000}], test: ftest[:{200}] }) # 定义 tokenize 函数 def tokenize_fn(examples): return self.tokenizer( examples[text], paddingmax_length, truncationTrue, max_lengthself.config.max_length, ) # 应用 tokenize train_dataset raw_datasets[train].map(tokenize_fn, batchedTrue) eval_dataset raw_datasets[test].map(tokenize_fn, batchedTrue) # 移除不需要的列重命名标签列 train_dataset train_dataset.remove_columns([text]) eval_dataset eval_dataset.remove_columns([text]) train_dataset train_dataset.rename_column(label, labels) eval_dataset eval_dataset.rename_column(label, labels) # 设置格式为 PyTorch tensor train_dataset.set_format(torch) eval_dataset.set_format(torch) return train_dataset, eval_dataset # # 模型管理模块 # class ModelManager: 模型保存和加载管理器 def __init__(self, config: FineTuneConfig): self.config config def save_model( self, model: AutoModelForSequenceClassification, tokenizer: AutoTokenizer, extra_info: Optional[Dict] None, ): 完整保存微调后的模型 保存内容 1. 模型权重 (model.safetensors) 2. 模型配置 (config.json) 3. Tokenizer 文件 (vocab.txt, tokenizer_config.json, etc.) 4. 额外元信息 (training_info.json) save_dir self.config.save_dir os.makedirs(save_dir, exist_okTrue) # 保存模型权重 配置 model.save_pretrained( save_dir, safe_serializationTrue, # 使用 safetensors 格式 ) # 保存 tokenizer tokenizer.save_pretrained(save_dir) # 保存训练元信息 training_info { model_name: self.config.model_name, num_labels: self.config.num_labels, max_length: self.config.max_length, learning_rate: self.config.learning_rate, num_epochs: self.config.num_epochs, save_time: datetime.now().isoformat(), transformers_version: __import__(transformers).__version__, torch_version: torch.__version__, } if extra_info: training_info.update(extra_info) with open(os.path.join(save_dir, training_info.json), w) as f: json.dump(training_info, f, indent2, ensure_asciiFalse) # 验证保存的文件 saved_files os.listdir(save_dir) print(f\n模型已保存到: {save_dir}) print(f保存的文件列表:) for f in sorted(saved_files): file_path os.path.join(save_dir, f) file_size os.path.getsize(file_path) print(f - {f} ({file_size / 1024:.1f} KB)) return save_dir def load_model(self, save_dir: str) - Tuple[AutoModelForSequenceClassification, AutoTokenizer, Dict]: 完整加载微调后的模型 加载内容 1. 模型权重和配置 2. Tokenizer 3. 训练元信息 # 加载 tokenizer tokenizer AutoTokenizer.from_pretrained(save_dir) # 加载模型自动读取 config.json model AutoModelForSequenceClassification.from_pretrained(save_dir) # 设置为评估模式 model.eval() # 移到 GPU如果可用 device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) # 加载训练元信息 info_path os.path.join(save_dir, training_info.json) training_info {} if os.path.exists(info_path): with open(info_path, r) as f: training_info json.load(f) print(f\n模型已从 {save_dir} 加载) print(f设备: {device}) if training_info: print(f原始训练信息:) for k, v in training_info.items(): print(f {k}: {v}) return model, tokenizer, training_info def verify_model_consistency( self, original_model: AutoModelForSequenceClassification, loaded_model: AutoModelForSequenceClassification, ) - bool: 验证保存前后的模型权重是否一致 original_state original_model.state_dict() loaded_state loaded_model.state_dict() if set(original_state.keys()) ! set(loaded_state.keys()): print(警告模型参数键不匹配) return False all_match True for key in original_state: if not torch.equal(original_state[key], loaded_state[key]): print(f参数 {key} 不匹配) all_match False if all_match: print(验证通过保存前后的模型权重完全一致。) return all_match # # 推理模块 # class ModelInference: 模型推理模块 def __init__(self, model, tokenizer, devicecpu): self.model model self.tokenizer tokenizer self.device device self.id2label model.config.id2label def predict(self, texts: List[str]) - List[Dict]: 对文本列表进行预测 # 编码输入 inputs self.tokenizer( texts, paddingTrue, truncationTrue, max_length128, return_tensorspt ).to(self.device) # 推理 self.model.eval() with torch.no_grad(): outputs self.model(**inputs) probabilities torch.nn.functional.softmax(outputs.logits, dim-1) predictions torch.argmax(probabilities, dim-1) # 组织结果 results [] for i, text in enumerate(texts): result { text: text, predicted_label: predictions[i].item(), predicted_label_name: self.id2label.get(predictions[i].item(), str(predictions[i].item())), probabilities: { self.id2label.get(j, str(j)): probabilities[i][j].item() for j in range(probabilities.shape[1]) }, } results.append(result) return results def predict_single(self, text: str) - Dict: 对单条文本进行预测 return self.predict([text])[0] # # 主流程 # def main(): 主流程微调 - 保存 - 加载 - 验证 - 推理 config FineTuneConfig() # 设置随机种子 set_seed(config.seed) random.seed(config.seed) np.random.seed(config.seed) torch.manual_seed(config.seed) print( * 60) print(步骤 1: 准备数据) print( * 60) data_prep DataPreparation(config) train_dataset, eval_dataset data_prep.prepare_data() tokenizer data_prep.tokenizer print(f训练集大小: {len(train_dataset)}) print(f评估集大小: {len(eval_dataset)}) print(\n * 60) print(步骤 2: 加载模型并微调) print( * 60) model AutoModelForSequenceClassification.from_pretrained( config.model_name, num_labelsconfig.num_labels, ) # 定义评估指标 def compute_metrics(eval_pred): logits, labels eval_pred predictions np.argmax(logits, axis-1) accuracy (predictions labels).mean() return {accuracy: accuracy} # 训练参数 training_args TrainingArguments( output_dir./training_output, num_train_epochsconfig.num_epochs, per_device_train_batch_sizeconfig.batch_size, per_device_eval_batch_sizeconfig.eval_batch_size, learning_rateconfig.learning_rate, warmup_ratioconfig.warmup_ratio, weight_decayconfig.weight_decay, logging_steps50, evaluation_strategyepoch, save_strategyno, # 我们手动保存 seedconfig.seed, report_tonone, ) # 创建 Trainer trainer Trainer( modelmodel, argstraining_args, train_datasettrain_dataset, eval_dataseteval_dataset, compute_metricscompute_metrics, ) # 训练 train_result trainer.train() print(f训练完成: {train_result}) # 评估 eval_result trainer.evaluate() print(f评估结果: {eval_result}) print(\n * 60) print(步骤 3: 保存模型) print( * 60) manager ModelManager(config) save_dir manager.save_model( modelmodel, tokenizertokenizer, extra_info{ train_loss: train_result.training_loss, eval_accuracy: eval_result.get(eval_accuracy, 0), } ) print(\n * 60) print(步骤 4: 重新加载模型) print( * 60) loaded_model, loaded_tokenizer, training_info manager.load_model(save_dir) print(\n * 60) print(步骤 5: 验证模型一致性) print( * 60) # 将原始模型也移到相同设备 device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) manager.verify_model_consistency(model, loaded_model) print(\n * 60) print(步骤 6: 使用加载的模型进行推理) print( * 60) inference ModelInference(loaded_model, loaded_tokenizer, device) test_texts [ This is the best movie I have ever seen! The acting was superb., Absolutely terrible. The plot made no sense and the acting was wooden., It was an okay movie. Some good parts but overall forgettable., A masterpiece of modern cinema. Every scene was beautifully crafted., Dont waste your money. The worst film of the year., ] results inference.predict(test_texts) print(\n推理结果:) print(- * 60) for result in results: print(f文本: {result[text][:60]}...) print(f预测标签: {result[predicted_label_name]}) print(f概率分布: {result[probabilities]}) print(- * 60) print(\n * 60) print(完整流程执行完毕) print( * 60) if __name__ __main__: main()常见陷阱与注意事项1. 忘记保存 Tokenizer这是最常见的错误。即使模型权重保存正确如果 tokenizer 不匹配预测结果也会完全错误。# 错误只保存模型 model.save_pretrained(./my_model) # 忘记保存 tokenizer # 正确同时保存 model.save_pretrained(./my_model) tokenizer.save_pretrained(./my_model) # 必须保存2. 加载后忘记设置 eval 模式# 错误 model AutoModelForSequenceClassification.from_pretrained(./my_model) # 直接预测dropout 仍然活跃结果不确定 # 正确 model AutoModelForSequenceClassification.from_pretrained(./my_model) model.eval() # 必须设置3. 跨版本兼容性问题不同版本的transformers库可能有不兼容的变更# 保存训练时的版本信息 import transformers import torch training_info { transformers_version: transformers.__version__, torch_version: torch.__version__, } # 加载时检查版本兼容性4. 自定义模型类的保存如果修改了模型架构如添加自定义头需要确保自定义类在加载时可被导入# 自定义模型类必须定义在可导入的模块中 class CustomBertModel(BertPreTrainedModel): def __init__(self, config): super().__init__(config) self.bert BertModel(config) self.custom_head nn.Linear(config.hidden_size, config.num_labels) def forward(self, input_ids, attention_maskNone): outputs self.bert(input_ids, attention_maskattention_mask) return self.custom_head(outputs.pooler_output) # 保存时需要指定自定义类 model.save_pretrained(./custom_model) # 加载时需要指定自定义类 from my_module import CustomBertModel # 必须可导入 model CustomBertModel.from_pretrained(./custom_model)5. 大模型分片保存对于大模型使用max_shard_size参数进行分片model.save_pretrained( ./large_model, max_shard_size1GB, # 每个分片最大 1GB )6. 半精度模型保存如果使用混合精度训练保存时注意精度# 保存 FP16 模型 model.half() model.save_pretrained(./fp16_model) # 加载时指定精度 model AutoModelForSequenceClassification.from_pretrained( ./fp16_model, torch_dtypetorch.float16, )7. 安全性考虑避免使用torch.save(model, path)保存整个模型对象因为pickle反序列化存在安全风险。始终使用save_pretrained()方法。总结正确保存和加载 Hugging Face 微调模型的关键要点始终使用save_pretrained()和from_pretrained()这是官方推荐的方式会自动保存和加载所有必要组件。同时保存模型和 tokenizer两者缺一不可必须保存在同一目录中。加载后设置 eval 模式model.eval()确保 dropout 等层在推理时行为正确。使用 safetensors 格式比 pickle 更安全、更快是现代推荐做法。保存训练元信息记录版本号、超参数等信息便于复现和调试。验证模型一致性加载后检查权重是否与保存前一致确保保存过程无损。避免直接torch.save(model)这种方式存在安全风险和跨环境兼容性问题。断点续训需保存完整状态包括 optimizer、scheduler、随机状态等。通过遵循这些最佳实践可以确保微调后的 Transformer 模型被正确保存和加载避免常见的坑和陷阱。无论是用于生产部署还是学术研究正确的模型持久化都是确保结果可复现的基础。 【Bug已解决】Saving and reload huggingface fine-tuned transformer 解决方案
【Bug已解决】Saving and reload huggingface fine-tuned transformer 解决方案问题描述在使用 Hugging Face Transformers 库对预训练模型进行微调fine-tuning后许多开发者会遇到一个常见且令人头疼的问题如何正确地保存微调后的模型并在后续使用时完整地重新加载它。这个看似简单的操作实际上涉及多个层面的复杂性包括模型权重、tokenizer、配置文件、训练状态等的持久化。典型的报错场景包括保存后重新加载时模型输出全是NaN或预测结果完全随机加载时报错AttributeError: XXX object has no attribute XXXTokenizer 与模型不匹配导致输入处理不一致使用torch.save()保存后跨环境加载出现ModuleNotFoundError保存的模型文件过大或过小怀疑保存不完整这些问题的根本原因在于 Hugging Face 模型的保存机制与普通 PyTorch 模型有所不同需要同时保存多个组件才能确保模型的完整性和可复现性。错误复现场景一仅保存模型权重导致加载失败# 错误做法只保存了 state_dict import torch from transformers import AutoModelForSequenceClassification, AutoTokenizer from torch.utils.data import DataLoader # 加载预训练模型 model AutoModelForSequenceClassification.from_pretrained( bert-base-uncased, num_labels2 ) tokenizer AutoTokenizer.from_pretrained(bert-base-uncased) # 模拟微调过程简化版 optimizer torch.optim.AdamW(model.parameters(), lr2e-5) model.train() for epoch in range(3): for batch in dataloader: outputs model(**batch) loss outputs.loss loss.backward() optimizer.step() optimizer.zero_grad() # 错误仅保存 state_dict torch.save(model.state_dict(), ./fine_tuned_model.pt) # 后续加载时出现问题 model_loaded AutoModelForSequenceClassification.from_pretrained( bert-base-uncased, num_labels2 ) model_loaded.load_state_dict(torch.load(./fine_tuned_model.pt)) # 问题tokenizer 没有保存配置可能不匹配 # 在新环境中可能报错或结果异常场景二跨环境加载报错# 在环境A中保存 torch.save(model, ./full_model.pt) # 在环境B中加载 model torch.load(./full_model.pt) # 报错ModuleNotFoundError: No module named transformers.models.bert.modeling_bert # 或AttributeError: Cannot get attribute XXX on module transformers...场景三保存后预测结果异常# 保存模型 model.save_pretrained(./my_model) tokenizer.save_pretrained(./my_model) # 加载模型 model AutoModelForSequenceClassification.from_pretrained(./my_model) # 但忘记加载 tokenizer使用了默认 tokenizer tokenizer AutoTokenizer.from_pretrained(bert-base-uncased) # 结果预测结果完全错误因为 tokenizer 配置可能不一致根因分析1. Hugging Face 模型的组件结构Hugging Face 的 Transformer 模型不是一个简单的nn.Module它由多个组件构成模型权重model weightspytorch_model.bin或model.safetensors模型配置config.json包含模型架构、超参数等信息Tokenizer 文件vocab.txt、tokenizer_config.json、special_tokens_map.json等生成配置generation_config.json对于生成模型训练状态optimizer state、scheduler state、random state 等仅保存state_dict会丢失配置信息导致加载时使用默认配置可能与微调时的配置不一致。2.torch.save()的序列化问题torch.save()使用 Python 的pickle模块进行序列化。当直接保存整个模型对象时# 这种方式保存的是模型的类路径和实例属性 torch.save(model, model.pt) # pickle 会记录: transformers.models.bert.modeling_bert.BertForSequenceClassification # 在新环境中如果版本不同类路径可能变化导致反序列化失败3. Tokenizer 的重要性Tokenizer 决定了文本如何被切分和编码。微调过程中如果修改了 tokenizer如添加特殊 token不保存 tokenizer 会导致加载后使用原始 tokenizer输入编码不一致。4. 模型评估模式未设置加载模型后忘记调用model.eval()导致 dropout 等层仍处于训练模式预测结果不确定。解决方案方案一使用save_pretrained()和from_pretrained()推荐这是 Hugging Face 官方推荐的方式会保存所有必要组件import torch from transformers import ( AutoModelForSequenceClassification, AutoTokenizer, TrainingArguments, Trainer, AutoConfig ) from datasets import load_dataset import numpy as np import os # # 完整的微调 保存 加载流程 # def fine_tune_and_save(): 完整的微调、保存和加载流程 model_name bert-base-uncased save_directory ./fine_tuned_bert # 1. 加载 tokenizer 和模型 tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForSequenceClassification.from_pretrained( model_name, num_labels2, # 自定义配置 hidden_dropout_prob0.1, attention_probs_dropout_prob0.1, ) print(f模型配置: {model.config}) print(f模型参数量: {sum(p.numel() for p in model.parameters()):,}) # 2. 准备数据集 dataset load_dataset(imdb, split{train: train[:1000], test: test[:200]}) def tokenize_function(examples): return tokenizer( examples[text], paddingmax_length, truncationTrue, max_length128, ) tokenized_datasets dataset.map(tokenize_function, batchedTrue) tokenized_datasets tokenized_datasets.remove_columns([text]) tokenized_datasets tokenized_datasets.rename_column(label, labels) tokenized_datasets.set_format(torch) # 3. 训练配置 training_args TrainingArguments( output_dir./training_output, num_train_epochs3, per_device_train_batch_size8, per_device_eval_batch_size16, warmup_steps100, weight_decay0.01, logging_dir./logs, logging_steps50, evaluation_strategyepoch, save_strategyepoch, load_best_model_at_endTrue, metric_for_best_modelaccuracy, ) # 4. 定义评估指标 def compute_metrics(eval_pred): logits, labels eval_pred predictions np.argmax(logits, axis-1) accuracy (predictions labels).mean() return {accuracy: accuracy} # 5. 创建 Trainer 并训练 trainer Trainer( modelmodel, argstraining_args, train_datasettokenized_datasets[train], eval_datasettokenized_datasets[test], compute_metricscompute_metrics, ) trainer.train() # 6. 评估 eval_results trainer.evaluate() print(f评估结果: {eval_results}) # # 关键正确保存模型和 tokenizer # # 创建保存目录 os.makedirs(save_directory, exist_okTrue) # 保存模型包含权重和配置 model.save_pretrained(save_directory) # 保存 tokenizer包含词表和配置 tokenizer.save_pretrained(save_directory) # 验证保存的文件 saved_files os.listdir(save_directory) print(f保存的文件: {saved_files}) # 预期文件列表: # - config.json (模型配置) # - model.safetensors 或 pytorch_model.bin (模型权重) # - vocab.txt (词表) # - tokenizer_config.json (tokenizer配置) # - special_tokens_map.json (特殊token映射) # - tokenizer.json (fast tokenizer文件) return save_directory def load_and_predict(save_directory): 正确加载模型并进行预测 # 1. 加载 tokenizer必须从同一目录加载 tokenizer AutoTokenizer.from_pretrained(save_directory) # 2. 加载模型自动读取 config.json model AutoModelForSequenceClassification.from_pretrained(save_directory) # 3. 设置为评估模式关键 model.eval() # 4. 如果有 GPU移到 GPU device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) # 5. 进行预测 test_texts [ This movie is absolutely fantastic! Highly recommended., Terrible film. Complete waste of time and money., An average movie, nothing special but not bad either., ] # 编码输入 inputs tokenizer( test_texts, paddingTrue, truncationTrue, max_length128, return_tensorspt ).to(device) # 推理 with torch.no_grad(): outputs model(**inputs) predictions torch.nn.functional.softmax(outputs.logits, dim-1) predicted_labels torch.argmax(predictions, dim-1) # 输出结果 labels [负面, 正面] for text, pred, probs in zip(test_texts, predicted_labels, predictions): print(f文本: {text[:50]}...) print(f预测: {labels[pred.item()]} (置信度: {probs[pred].item():.4f})) print() # 执行完整流程 if __name__ __main__: save_dir fine_tune_and_save() load_and_predict(save_dir)方案二保存完整训练状态用于断点续训如果需要中断后继续训练需要保存更多状态import torch from transformers import AutoModelForSequenceClassification, AutoTokenizer import os import json from datetime import datetime def save_training_checkpoint(model, tokenizer, optimizer, scheduler, epoch, step, save_dir): 保存完整的训练检查点支持断点续训 os.makedirs(save_dir, exist_okTrue) # 1. 保存模型和 tokenizer使用 Hugging Face 方法 model.save_pretrained(save_dir) tokenizer.save_pretrained(save_dir) # 2. 保存训练状态使用 PyTorch 方法 checkpoint { epoch: epoch, step: step, optimizer_state_dict: optimizer.state_dict(), scheduler_state_dict: scheduler.state_dict() if scheduler else None, random_state: { python: random.getstate(), numpy: np.random.get_state(), torch: torch.get_rng_state(), cuda: torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None, }, timestamp: datetime.now().isoformat(), } torch.save(checkpoint, os.path.join(save_dir, training_state.pt)) # 3. 保存训练元信息 meta_info { epoch: epoch, step: step, model_type: model.__class__.__name__, save_time: datetime.now().isoformat(), } with open(os.path.join(save_dir, training_meta.json), w) as f: json.dump(meta_info, f, indent2) print(f训练检查点已保存到 {save_dir}) print(f - Epoch: {epoch}, Step: {step}) def load_training_checkpoint(save_dir, model_nameNone): 加载训练检查点恢复训练状态 # 1. 加载模型和 tokenizer model AutoModelForSequenceClassification.from_pretrained(save_dir) tokenizer AutoTokenizer.from_pretrained(save_dir) # 2. 加载训练状态 checkpoint torch.load(os.path.join(save_dir, training_state.pt)) # 3. 恢复随机状态确保可复现性 random.setstate(checkpoint[random_state][python]) np.random.set_state(checkpoint[random_state][numpy]) torch.set_rng_state(checkpoint[random_state][torch]) if checkpoint[random_state][cuda] is not None: torch.cuda.set_rng_state_all(checkpoint[random_state][cuda]) # 4. 创建 optimizer 并恢复状态 optimizer torch.optim.AdamW(model.parameters(), lr2e-5) optimizer.load_state_dict(checkpoint[optimizer_state_dict]) # 5. 恢复 scheduler如果有 scheduler None if checkpoint[scheduler_state_dict] is not None: from transformers import get_linear_schedule_with_warmup scheduler get_linear_schedule_with_warmup( optimizer, num_warmup_steps100, num_training_steps10000 ) scheduler.load_state_dict(checkpoint[scheduler_state_dict]) print(f训练检查点已从 {save_dir} 加载) print(f - 恢复到 Epoch: {checkpoint[epoch]}, Step: {checkpoint[step]})  return model, tokenizer, optimizer, scheduler, checkpoint[epoch], checkpoint[step]方案三使用 Safetensors 格式更安全from transformers import AutoModelForSequenceClassification # 保存为 safetensors 格式默认行为更安全 model AutoModelForSequenceClassification.from_pretrained(bert-base-uncased) # safetensors 格式避免了 pickle 的安全风险 # 且加载速度更快 model.save_pretrained( ./safe_model, safe_serializationTrue, # 使用 safetensors 格式 max_shard_size500MB, # 大模型分片保存 ) # 加载时自动识别格式 model AutoModelForSequenceClassification.from_pretrained(./safe_model)完整修复代码下面是一个完整的、可运行的端到端示例涵盖微调、保存、加载和验证 完整的 Hugging Face Transformer 微调模型保存与加载示例 包含数据准备、微调训练、模型保存、模型加载、结果验证 import torch import torch.nn as nn import numpy as np import random import os import json from datetime import datetime from typing import Dict, List, Tuple, Optional # Hugging Face 相关导入 from transformers import ( AutoModelForSequenceClassification, AutoTokenizer, AutoConfig, TrainingArguments, Trainer, get_linear_schedule_with_warmup, set_seed, ) from datasets import load_dataset, Dataset from torch.utils.data import DataLoader # # 配置类 # class FineTuneConfig: 微调配置 model_name: str bert-base-uncased num_labels: int 2 max_length: int 128 batch_size: int 8 eval_batch_size: int 16 learning_rate: float 2e-5 num_epochs: int 3 warmup_ratio: float 0.1 weight_decay: float 0.01 seed: int 42 save_dir: str ./fine_tuned_model_complete classmethod def to_dict(cls): return {k: v for k, v in cls.__dict__.items() if not k.startswith(_)} # # 数据准备模块 # class DataPreparation: 数据准备和预处理 def __init__(self, config: FineTuneConfig): self.config config self.tokenizer AutoTokenizer.from_pretrained(config.model_name) def prepare_data(self) - Tuple[Dataset, Dataset]: 准备训练和评估数据集 # 使用 IMDB 数据集的子集 raw_datasets load_dataset(imdb, split{ train: ftrain[:{1000}], test: ftest[:{200}] }) # 定义 tokenize 函数 def tokenize_fn(examples): return self.tokenizer( examples[text], paddingmax_length, truncationTrue, max_lengthself.config.max_length, ) # 应用 tokenize train_dataset raw_datasets[train].map(tokenize_fn, batchedTrue) eval_dataset raw_datasets[test].map(tokenize_fn, batchedTrue) # 移除不需要的列重命名标签列 train_dataset train_dataset.remove_columns([text]) eval_dataset eval_dataset.remove_columns([text]) train_dataset train_dataset.rename_column(label, labels) eval_dataset eval_dataset.rename_column(label, labels) # 设置格式为 PyTorch tensor train_dataset.set_format(torch) eval_dataset.set_format(torch) return train_dataset, eval_dataset # # 模型管理模块 # class ModelManager: 模型保存和加载管理器 def __init__(self, config: FineTuneConfig): self.config config def save_model( self, model: AutoModelForSequenceClassification, tokenizer: AutoTokenizer, extra_info: Optional[Dict] None, ): 完整保存微调后的模型 保存内容 1. 模型权重 (model.safetensors) 2. 模型配置 (config.json) 3. Tokenizer 文件 (vocab.txt, tokenizer_config.json, etc.) 4. 额外元信息 (training_info.json) save_dir self.config.save_dir os.makedirs(save_dir, exist_okTrue) # 保存模型权重 配置 model.save_pretrained( save_dir, safe_serializationTrue, # 使用 safetensors 格式 ) # 保存 tokenizer tokenizer.save_pretrained(save_dir) # 保存训练元信息 training_info { model_name: self.config.model_name, num_labels: self.config.num_labels, max_length: self.config.max_length, learning_rate: self.config.learning_rate, num_epochs: self.config.num_epochs, save_time: datetime.now().isoformat(), transformers_version: __import__(transformers).__version__, torch_version: torch.__version__, } if extra_info: training_info.update(extra_info) with open(os.path.join(save_dir, training_info.json), w) as f: json.dump(training_info, f, indent2, ensure_asciiFalse) # 验证保存的文件 saved_files os.listdir(save_dir) print(f\n模型已保存到: {save_dir}) print(f保存的文件列表:) for f in sorted(saved_files): file_path os.path.join(save_dir, f) file_size os.path.getsize(file_path) print(f - {f} ({file_size / 1024:.1f} KB)) return save_dir def load_model(self, save_dir: str) - Tuple[AutoModelForSequenceClassification, AutoTokenizer, Dict]: 完整加载微调后的模型 加载内容 1. 模型权重和配置 2. Tokenizer 3. 训练元信息 # 加载 tokenizer tokenizer AutoTokenizer.from_pretrained(save_dir) # 加载模型自动读取 config.json model AutoModelForSequenceClassification.from_pretrained(save_dir) # 设置为评估模式 model.eval() # 移到 GPU如果可用 device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) # 加载训练元信息 info_path os.path.join(save_dir, training_info.json) training_info {} if os.path.exists(info_path): with open(info_path, r) as f: training_info json.load(f) print(f\n模型已从 {save_dir} 加载) print(f设备: {device}) if training_info: print(f原始训练信息:) for k, v in training_info.items(): print(f {k}: {v}) return model, tokenizer, training_info def verify_model_consistency( self, original_model: AutoModelForSequenceClassification, loaded_model: AutoModelForSequenceClassification, ) - bool: 验证保存前后的模型权重是否一致 original_state original_model.state_dict() loaded_state loaded_model.state_dict() if set(original_state.keys()) ! set(loaded_state.keys()): print(警告模型参数键不匹配) return False all_match True for key in original_state: if not torch.equal(original_state[key], loaded_state[key]): print(f参数 {key} 不匹配) all_match False if all_match: print(验证通过保存前后的模型权重完全一致。) return all_match # # 推理模块 # class ModelInference: 模型推理模块 def __init__(self, model, tokenizer, devicecpu): self.model model self.tokenizer tokenizer self.device device self.id2label model.config.id2label def predict(self, texts: List[str]) - List[Dict]: 对文本列表进行预测 # 编码输入 inputs self.tokenizer( texts, paddingTrue, truncationTrue, max_length128, return_tensorspt ).to(self.device) # 推理 self.model.eval() with torch.no_grad(): outputs self.model(**inputs) probabilities torch.nn.functional.softmax(outputs.logits, dim-1) predictions torch.argmax(probabilities, dim-1) # 组织结果 results [] for i, text in enumerate(texts): result { text: text, predicted_label: predictions[i].item(), predicted_label_name: self.id2label.get(predictions[i].item(), str(predictions[i].item())), probabilities: { self.id2label.get(j, str(j)): probabilities[i][j].item() for j in range(probabilities.shape[1]) }, } results.append(result) return results def predict_single(self, text: str) - Dict: 对单条文本进行预测 return self.predict([text])[0] # # 主流程 # def main(): 主流程微调 - 保存 - 加载 - 验证 - 推理 config FineTuneConfig() # 设置随机种子 set_seed(config.seed) random.seed(config.seed) np.random.seed(config.seed) torch.manual_seed(config.seed) print( * 60) print(步骤 1: 准备数据) print( * 60) data_prep DataPreparation(config) train_dataset, eval_dataset data_prep.prepare_data() tokenizer data_prep.tokenizer print(f训练集大小: {len(train_dataset)}) print(f评估集大小: {len(eval_dataset)}) print(\n * 60) print(步骤 2: 加载模型并微调) print( * 60) model AutoModelForSequenceClassification.from_pretrained( config.model_name, num_labelsconfig.num_labels, ) # 定义评估指标 def compute_metrics(eval_pred): logits, labels eval_pred predictions np.argmax(logits, axis-1) accuracy (predictions labels).mean() return {accuracy: accuracy} # 训练参数 training_args TrainingArguments( output_dir./training_output, num_train_epochsconfig.num_epochs, per_device_train_batch_sizeconfig.batch_size, per_device_eval_batch_sizeconfig.eval_batch_size, learning_rateconfig.learning_rate, warmup_ratioconfig.warmup_ratio, weight_decayconfig.weight_decay, logging_steps50, evaluation_strategyepoch, save_strategyno, # 我们手动保存 seedconfig.seed, report_tonone, ) # 创建 Trainer trainer Trainer( modelmodel, argstraining_args, train_datasettrain_dataset, eval_dataseteval_dataset, compute_metricscompute_metrics, ) # 训练 train_result trainer.train() print(f训练完成: {train_result}) # 评估 eval_result trainer.evaluate() print(f评估结果: {eval_result}) print(\n * 60) print(步骤 3: 保存模型) print( * 60) manager ModelManager(config) save_dir manager.save_model( modelmodel, tokenizertokenizer, extra_info{ train_loss: train_result.training_loss, eval_accuracy: eval_result.get(eval_accuracy, 0), } ) print(\n * 60) print(步骤 4: 重新加载模型) print( * 60) loaded_model, loaded_tokenizer, training_info manager.load_model(save_dir) print(\n * 60) print(步骤 5: 验证模型一致性) print( * 60) # 将原始模型也移到相同设备 device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) manager.verify_model_consistency(model, loaded_model) print(\n * 60) print(步骤 6: 使用加载的模型进行推理) print( * 60) inference ModelInference(loaded_model, loaded_tokenizer, device) test_texts [ This is the best movie I have ever seen! The acting was superb., Absolutely terrible. The plot made no sense and the acting was wooden., It was an okay movie. Some good parts but overall forgettable., A masterpiece of modern cinema. Every scene was beautifully crafted., Dont waste your money. The worst film of the year., ] results inference.predict(test_texts) print(\n推理结果:) print(- * 60) for result in results: print(f文本: {result[text][:60]}...) print(f预测标签: {result[predicted_label_name]}) print(f概率分布: {result[probabilities]}) print(- * 60) print(\n * 60) print(完整流程执行完毕) print( * 60) if __name__ __main__: main()常见陷阱与注意事项1. 忘记保存 Tokenizer这是最常见的错误。即使模型权重保存正确如果 tokenizer 不匹配预测结果也会完全错误。# 错误只保存模型 model.save_pretrained(./my_model) # 忘记保存 tokenizer # 正确同时保存 model.save_pretrained(./my_model) tokenizer.save_pretrained(./my_model) # 必须保存2. 加载后忘记设置 eval 模式# 错误 model AutoModelForSequenceClassification.from_pretrained(./my_model) # 直接预测dropout 仍然活跃结果不确定 # 正确 model AutoModelForSequenceClassification.from_pretrained(./my_model) model.eval() # 必须设置3. 跨版本兼容性问题不同版本的transformers库可能有不兼容的变更# 保存训练时的版本信息 import transformers import torch training_info { transformers_version: transformers.__version__, torch_version: torch.__version__, } # 加载时检查版本兼容性4. 自定义模型类的保存如果修改了模型架构如添加自定义头需要确保自定义类在加载时可被导入# 自定义模型类必须定义在可导入的模块中 class CustomBertModel(BertPreTrainedModel): def __init__(self, config): super().__init__(config) self.bert BertModel(config) self.custom_head nn.Linear(config.hidden_size, config.num_labels) def forward(self, input_ids, attention_maskNone): outputs self.bert(input_ids, attention_maskattention_mask) return self.custom_head(outputs.pooler_output) # 保存时需要指定自定义类 model.save_pretrained(./custom_model) # 加载时需要指定自定义类 from my_module import CustomBertModel # 必须可导入 model CustomBertModel.from_pretrained(./custom_model)5. 大模型分片保存对于大模型使用max_shard_size参数进行分片model.save_pretrained( ./large_model, max_shard_size1GB, # 每个分片最大 1GB )6. 半精度模型保存如果使用混合精度训练保存时注意精度# 保存 FP16 模型 model.half() model.save_pretrained(./fp16_model) # 加载时指定精度 model AutoModelForSequenceClassification.from_pretrained( ./fp16_model, torch_dtypetorch.float16, )7. 安全性考虑避免使用torch.save(model, path)保存整个模型对象因为pickle反序列化存在安全风险。始终使用save_pretrained()方法。总结正确保存和加载 Hugging Face 微调模型的关键要点始终使用save_pretrained()和from_pretrained()这是官方推荐的方式会自动保存和加载所有必要组件。同时保存模型和 tokenizer两者缺一不可必须保存在同一目录中。加载后设置 eval 模式model.eval()确保 dropout 等层在推理时行为正确。使用 safetensors 格式比 pickle 更安全、更快是现代推荐做法。保存训练元信息记录版本号、超参数等信息便于复现和调试。验证模型一致性加载后检查权重是否与保存前一致确保保存过程无损。避免直接torch.save(model)这种方式存在安全风险和跨环境兼容性问题。断点续训需保存完整状态包括 optimizer、scheduler、随机状态等。通过遵循这些最佳实践可以确保微调后的 Transformer 模型被正确保存和加载避免常见的坑和陷阱。无论是用于生产部署还是学术研究正确的模型持久化都是确保结果可复现的基础。