pytorch 神经网络预训练模型的保存与提取

1、建立神经网络、训练神经网络

神经网络的建立与训练代码，我们会在pytorch手写数字识别篇来分享具体的代码含义

2、保存神经网络

pytorch的神经网络保存一共有2种方式

方式一：

# 仅保存CNN参数，速度较快

torch.save(cnn.state_dict(), './model/CNN_NO.pkl')

方式二：

# 保存CNN整个结构，速度较慢

torch.save(cnn(), './model/CNN.pkl')

3、神经网络的提取

pytorch既然有2种保存方式，必然有2种提取方式

方式一：

model = CNN() # 提取模型参数，速度较快

model.load_state_dict(torch.load('./model/CNN_state_dict.pkl'))

model.eval()

output= model(input) # 传入输入数据，进行预测

方式二：

model = CNN() # 提取整个模型，速度较慢

model = torch.load('./model/CNN.pkl'))

model.eval()

output= model(input) # 传入输入数据，进行预测

小编开发了小程序，可以体验不同的AI工具

微&*信搜索：启示AI科技

体验不一样的AI工具