torch.save

pytorch保存模型等相关参数,利用torch.save(),以及读取保存之后的文件

我对 torch.save 的理解是,这个函数可以将一个对象序列化后储存在一个文件中,然后再用torch.load加载回来。

所以你可以存这个训练好的模型model

torch.save(model, path)

然后加载回来后直接可以用

model = torch.load(path)
output = model(image)

好处是不需要知道之前的模型代码,直接加载调用即可,方便使用。但多存储了模型的结构数据,存起来肯能会慢一丢丢,到底会不会有太大影响,我也没测试过。

但从网上查到的都是说在训练的时候仅仅保存模型的state_dict即可,这样快,说是快,我也没对比过,之后有空了也许会测一下(挖坑)

我们也可以只存储模型的state_dict,就是模型的权重

torch.save(model.state_dict(), path)

载入权重时,要先用之前的模型初始化一个对象,然后使用load_state_dict函数来加载我们存储的权重

model = Model()
model.load_state_dict(torch.load(path))

使用样例 1

假设网络为model = Net(), optimizer = optim.Adam(model.parameters(), lr=args.lr), 假设在某个epoch,我们要保存模型参数,优化器参数以及epoch

1. 先建立一个字典,保存三个参数:

state = {‘net':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}

2.调用torch.save():

torch.save(state, dir)

其中dir表示保存文件的绝对路径+保存文件名,如'/home/qinying/Desktop/modelpara.pth'

当你想恢复某一阶段的训练(或者进行测试)时,那么就可以读取之前保存的网络模型参数等。

checkpoint = torch.load(dir)
model.load_state_dict(checkpoint['net'])
optimizer.load_state_dict(checkpoint['optimizer'])
start_epoch = checkpoint['epoch'] + 1

使用样例 2

https://github.com/Etherwave/minist

文章目录