模型的保存:
方式一:保存整个模型
model = resnet()
#-------省略训练过程------------
torch.save(model,path)
方式二:只保存模型参数
model = resnet()
#-------省略训练过程---------------
torch.save(model.state_dict,path)
模型的加载:
情况一:训练已经结束,测试时加载训练好的模型
1、在训练的py文件下面直接对模型进行测试:
model = resnet() # 定义模型
#----训练过程省略-----
torch.save(model,path) #保存模型
#----训练好的模型进行测试----
model.eval
#---加载测试集等步骤
2、单独写一个infer.py
①加载整个模型和参数,对应模型的保存方式一,即当时就保存了整个模型,此时直接加载,然后去测试;
model_path = './model_000009.pth'
model = torch.load(model_path) # 加载了整个模型
②只加载模型参数,对应模型保存的方式二,当时只保存了模型参数,此时需要先初始化模型,然后让模型加载训练好的模型的参数;
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
情况二:训练一半停了,想继续训练
只加载模型参数,网络结构从代码中创建
net = resnet() #代码中创建网络结构
path = './pth'
saved_state_dict = torch.load(path) #加载模型参数
net.load_state_dict(saved_state_dict) #应用到网络中
情况三:加载预训练模型
1、使用torchvision自带的预训练模型
#----------resnet系列------------
import torchvision.models as models
model = models.ResNet(pretrained=True)
model = models.resnet18(pretrained=True)
model = models.resnet34(pretrained=True)
model = models.resnet50(pretrained=True)
#---------VGG系列-------------
model = models.VGG(pretrained=True)
model = models.vgg11(pretrained=True)
model = models.vgg16(pretrained=True)
model = models.vgg16_bn(pretrained=True)
#---------alexnet-----------
Alexnet = models.alexnet()
#--------squeezenet--------
squeezenet = models.squeezenet1_0()
2、如果只需要加载预训练模型的部分参数&#