pytorch 保存和加载模型两种方法学习

一、保存整个神经网络的的结构信息和模型参数信息,save的对象是网络net

# 保存和加载整个模型
torch.save(model_object, 'resnet.pth')
model = torch.load('resnet.pth')

二、只保存神经网络的训练模型参数,save的对象是net.state_dict()

# 将my_resnet模型储存为my_resnet.pth
torch.save(my_resnet.state_dict(), "my_resnet.pth")
# 加载resnet,模型存放在my_resnet.pth
my_resnet.load_state_dict(torch.load("my_resnet.pth"))

其中my_resnet是my_resnet.pth对应的网络结构。
第一种方法相对更耗时。

猜你喜欢

转载自blog.csdn.net/qq_41872271/article/details/105367290
今日推荐