是的,你可以暂停训练并在程序暂停时保存权重。在PyTorch中,你可以使用torch.save()
函数来保存模型的权重。以下是一个示例代码:
import torch
import torchvision.models as models
# 创建模型
model = models.resnet18()
# 模型训练过程...
# 暂停训练并保存权重
torch.save(model.state_dict(), 'model_weights.pth')
在上述代码中,我们使用了torchvision.models
模块中的一个预训练模型作为示例。你可以根据自己的需求选择合适的模型。
torch.save()
函数接受两个参数:要保存的对象和保存的文件路径。在这个例子中,我们使用model.state_dict()
来保存模型的权重。你也可以选择保存整个模型,包括模型的结构和参数,使用torch.save(model, 'model.pth')
。
当你想要恢复训练时,可以使用torch.load()
函数加载保存的权重,并将其加载到模型中。以下是一个示例代码:
import torch
import torchvision.models as models
# 创建模型
model = models.resnet18()
# 加载保存的权重
model.load_state_dict(torch.load('model_weights.pth'))
# 恢复训练...
这样,你就可以在程序暂停时保存模型的权重,并在需要时恢复训练。
领取专属 10元无门槛券
手把手带您无忧上云