模型保存与加载

前言

先上代码

from transformers import Trainer，BertConfig，BertForSequenceClassification
import torch

......

# 创建Trainer
trainer = Trainer(
    model=model.to('cuda'),             
    args=training_args,                 
    train_dataset=train_dataset,        
    eval_dataset=test_dataset,          
)
# 开始训练
trainer.train()
# 开始评估模型
trainer.evaluate()
# 保存模型 会保存到配置的output_dir处
trainer.save_model()
torch.save(model.state_dict(), 'model_save.bin')

# 使用训练好的模型
output_config_file = './my_results/config.json'
output_model_file = 'model_save.bin'

config = BertConfig.from_json_file(output_config_file)
model = BertForSequenceClassification(config)
state_dict = torch.load(output_model_file)
model.load_state_dict(state_dict)

cache_dir="bert-base-uncased/"
tokenizer = BertTokenizer.from_pretrained(cache_dir)
data = tokenizer(['This is a good movie', 'This is a bad movie'], max_length=512, truncation=True, padding='max_length', return_tensors="pt")
print(model(**data))

Pytorch是深度学习领域中非常流行的框架之一，支持的模型保存格式包括.pt和.pth .bin .onnx。这几种格式的文件都可以保存Pytorch训练出的模型，但是它们的区别是什么呢？

模型的保存与加载到底在做什么？

保存训练好的模型参数，便于之后进行复现或继续训练

• model.state_dict()，模型每一层可学习的节点的参数，比如weight/bias；
• optimizer.state_dict()，模型的优化器中的参数；

.pt .pth .bin 以及 .onnx 是保存时约定的格式，对于torch.load()方法而言，不管你把后缀改成是什么，只要文件是对的都可以读取。
格式汇总：

来源：链接