13-pytorch加载上次训练结果文件后继续训练

最新推荐文章于 2024-07-19 16:16:46 发布

大鱼不吃鱼（wanna）

最新推荐文章于 2024-07-19 16:16:46 发布

阅读量982

点赞数 1

分类专栏： Python和机器学习文章标签： pytorch python 人工智能

本文链接：https://blog.csdn.net/qq_41286751/article/details/127137651

版权

Python和机器学习专栏收录该内容

24 篇文章 2 订阅

订阅专栏


# 如
model = model.restnet50()

# 如果有保存好的训练文件，在后面加上下面几句话
resume = 'checkpoint-480.pth'
checkpoint = torch.load(resume)
model.load_state_dict (checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])

保存模型的方法：

# 定义要保留的格式及数据
def save_checkpoint(path, model, optimizer):
    state = {
        'model': model.state_dict(),
        'optimizer': optimizer.state_dict()
    }
    torch.save(state, path)



# 在一个epoch结束后，写上：
save_checkpoint('checkpoint-%i.pth' % index, model, optimizer)



# optimizer举例
optimizer = torch.optim.SGD(
    model.parameters(),
    lr=LEARNING_RATE,
    momentum=MOMENTUM
)