2. 网络模型和权重的保存 wandb(weight and bias) in pytorch
一、网络模型或者权重的保存的方式
- 方式一: 只保存网络的权重参数
model = MainNet()
torch.save(model.state_dict(), "ckpts/params.pth")
- 方式二:保存网络和权重
model = MainNet()
torch.save(model, "ckpts/params.pth")
- 方式二点五:使用half()对精度进行缩减保存
model = MainNet()
torch.save(model.half(), "ckpts/params.pth")
- [注意]
- (1) 如果已经使用了方式一只保存了参数,也没有关系,重新在保存一次
net = MainNet()
net.load_state_dict(torch.load("ckpts/params.pth"))
torch.save(net.half(), 'ckpts/params_half.pth')
net = torch.load("ckpts/params_half.pth")
print(net)