如何在PyTorch中保存训练好的模型?我读到过:
Torch.save ()/torch.load()用于保存/加载可序列化对象。 model.state_dict()/model.load_state_dict()用于保存/加载模型状态。
如何在PyTorch中保存训练好的模型?我读到过:
Torch.save ()/torch.load()用于保存/加载可序列化对象。 model.state_dict()/model.load_state_dict()用于保存/加载模型状态。
当前回答
我总是喜欢使用Torch7 (.t7)或Pickle (.pth, .pt)来保存pytorch模型的权重。
其他回答
现在所有内容都写在官方教程中: https://pytorch.org/tutorials/beginner/saving_loading_models.html
关于如何保存和保存什么,您有几个选项,所有这些都在本教程中进行了解释。
这取决于你想做什么。
案例# 1:保存模型以便自己使用它进行推理:保存模型,恢复模型,然后将模型更改为评估模式。这样做是因为你通常有BatchNorm和Dropout图层,默认情况下在构造时处于火车模式:
torch.save(model.state_dict(), filepath)
#Later to restore:
model.load_state_dict(torch.load(filepath))
model.eval()
案例2:保存模型以稍后恢复训练:如果您需要继续训练即将保存的模型,那么您需要保存的不仅仅是模型。您还需要保存优化器的状态、epoch、分数等。你会这样做:
state = {
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
...
}
torch.save(state, filepath)
为了恢复训练,你可以这样做:state = torch.load(filepath),然后,恢复每个单独对象的状态,就像这样:
model.load_state_dict(state['state_dict'])
optimizer.load_state_dict(state['optimizer'])
由于您正在恢复训练,因此在加载时恢复状态后不要调用model.eval()。
案例# 3:模型被其他人使用,而不能访问您的代码: 在Tensorflow中,你可以创建一个.pb文件来定义模型的架构和权重。这非常方便,特别是在使用Tensorflow服务时。在Pytorch中这样做的等效方法是:
torch.save(model, filepath)
# Then later:
model = torch.load(filepath)
这种方式仍然不是防弹的,因为pytorch仍然在经历很多变化,我不推荐它。
我总是喜欢使用Torch7 (.t7)或Pickle (.pth, .pt)来保存pytorch模型的权重。
一个常见的PyTorch约定是使用.pt或.pth文件扩展名保存模型。
保存/加载整个模型
拯救策略:
path = "username/directory/lstmmodelgpu.pth"
torch.save(trainer, path)
负载:
(模型类必须在某处定义)
model.load_state_dict(torch.load(PATH))
model.eval()
我用这个方法,希望对大家有用。
num_labels = len(test_label_cols)
robertaclassificationtrain = '/dbfs/FileStore/tables/PM/TC/roberta_model'
robertaclassificationpath = "/dbfs/FileStore/tables/PM/TC/ROBERTACLASSIFICATION"
model = RobertaForSequenceClassification.from_pretrained(robertaclassificationpath,
num_labels=num_labels)
model.cuda()
model.load_state_dict(torch.load(robertaclassificationtrain))
model.eval()
我保存我的火车模型已经在“roberta_model”路径。保存一个火车模型。
torch.save(model.state_dict(), '/dbfs/FileStore/tables/PM/TC/roberta_model')