save_pretrained
vs save_model
的区别
在 Hugging Face transformers
库中,save_pretrained
和 save_model
都用于 保存模型,但它们的用途、适用范围和存储内容有所不同。
1. save_pretrained
适用于 Hugging Face 生态系统的模型保存
model.save_pretrained("./my_model")
tokenizer.save_pretrained("./my_model")
- 适用于
AutoModel
或Trainer
训练的 Hugging Face 模型 - 可以用
from_pretrained()
重新加载 - 保存
.bin
或.safetensors
权重 - 适用于 PyTorch 和 TensorFlow
- 推荐用于 Hugging Face 生态
2. save_model
适用于 Trainer
训练的模型
trainer.save_model("./my_model")
- 仅适用于
Trainer
训练的模型 - 通常用于在
Trainer
训练过程中保存 - 等价于
model.save_pretrained()
- 可以被
from_pretrained()
重新加载 - 适用于 PyTorch 和 TensorFlow
3. save_pretrained
vs save_model
主要区别
功能 | save_pretrained | save_model |
---|---|---|
适用范围 | 任何 Hugging Face 模型(AutoModel 、手动训练的 model ) | 仅 Trainer 训练的模型 |
适用于 Trainer 吗? | ✅ 适用 | ✅ 适用 |
能否用 from_pretrained() 重新加载 | ✅ 可以 | ✅ 可以 |
存储格式 | PyTorch (.bin )、TensorFlow (.ckpt ) | PyTorch (.bin )、TensorFlow (.ckpt ) |
是否保存 Tokenizer | ❌ 需要单独 tokenizer.save_pretrained() | ❌ 需要单独 tokenizer.save_pretrained() |
是否推荐 | ✅ 推荐 Hugging Face 用户使用 | 仅 Trainer 训练时使用 |
4. 什么时候用 save_pretrained
?
✅ 适用于
- 手动微调或加载 Hugging Face 预训练模型
- 无需
Trainer
也能使用 - 可以保存
tokenizer
model.save_pretrained("./my_model")
tokenizer.save_pretrained("./my_model")
# 重新加载
from transformers import AutoModel, AutoTokenizer
model = AutoModel.from_pretrained("./my_model")
tokenizer = AutoTokenizer.from_pretrained("./my_model")
5. 什么时候用 save_model
?
✅ 适用于
- 使用
Trainer
训练后保存模型 - 等价于
save_pretrained
- 如果
Trainer
训练,推荐trainer.save_model()
trainer.save_model("./my_model")
然后可以重新加载:
from transformers import AutoModel
model = AutoModel.from_pretrained("./my_model")
6. save_pretrained
vs save_model
结论
场景 | 推荐方法 |
---|---|
Hugging Face AutoModel 微调后保存 | ✅ model.save_pretrained() |
使用 Trainer 训练后保存 | ✅ trainer.save_model() |
加载 Hugging Face 预训练模型后保存 | ✅ model.save_pretrained() |
需要兼容 from_pretrained() 重新加载 | ✅ model.save_pretrained() |
✅ 结论:
- 推荐
save_pretrained
,更通用,适用于 Hugging Face 生态 save_model
仅适用于Trainer
训练的模型