pytorch加载训练好的模型用来测试或者处理

1.直接加载预训练模型

如果我们使用的模型和原模型完全一样，那么我们可以直接加载别人训练好的模型：

import torchvision.models as models
 
resnet50 = models.resnet50(pretrained=True)

如果只需要网络结构，不需要用预训练模型的参数来初始化，那么就是：

model =torchvision.models.resnet50(pretrained=False)

2.修改某一层

PyTorch中的torchvision里已经有很多常用的模型了，可以直接调用：

• AlexNet
• VGG
• ResNet
• SqueezeNet
• DenseNet

import torchvision.models as models
resnet18 = models.resnet18()
alexnet = models.alexnet()
squeezenet = models.squeezenet1_0()
densenet = models.densenet_161()

但是对于我们的任务而言有些层并不是直接能用，需要我们微微改一下，比如，resnet最后的全连接层是分1000类，而我们只有21类；又比如，resnet第一层卷积接收的通道是3， 我们可能输入图片的通道是4，那么可以通过以下方法修改：

resnet.conv1 = nn.Conv2d(4, 64, kernel_size=7, stride=2, padding=3, bias=False)
resnet.fc = nn.Linear(2048, 21)

3.加载部分预训练模型

其实大多数时候我们需要根据我们的任务调节我们的模型，所以很难保证模型和公开的模型完全一样，但是预训练模型的参数确实有助于提高训练的准确率，为了结合二者的优点，就需要我们加载部分预训练模型。

#加载model，model是自己定义好的模型
resnet50 = models.resnet50(pretrained=True) 
model =Net(...) 
 
#读取参数 
pretrained_dict =resnet50.state_dict() 
model_dict = model.state_dict() 
 
#将pretrained_dict里不属于model_dict的键剔除掉 
pretrained_dict =  {k: v for k, v in pretrained_dict.items() if k in model_dict} 
 
# 更新现有的model_dict 
model_dict.update(pretrained_dict) 

# 加载我们真正需要的state_dict
model.load_state_dict(model_dict)
 
# 加载我们真正需要的state_dict 
model.load_state_dict(model_dict)  

4. 保存和加载自己的模型

pytorch保存模型的方式有两种：
第一种：将整个网络都都保存下来
第二种：仅保存和加载模型参数（推荐使用这样的方法）

4.1 保存和加载整个模型

# 保存
torch.save(model_object, Path)

# 加载
model = torch.load(Path)

 4.2 仅保存和加载模型参数(推荐使用)
 

# ----------------保存模型参数--------------------------
torch.save(model.state_dict(), PATH)
 
#example
torch.save(resnet50.state_dict(),'ckp/model.pth')    

# ----------------加载模型参数--------------------------

model = ModelClass(*args, **kwargs) # 这是你后来设置的模型
model.load_state_dict(torch.load(PATH)) # 加载参数
 
#example
resnet=resnet50(pretrained=True)
resnet.load_state_dict(torch.load('ckp/model.pth'))

4.3 每个epoch保存一个模型参数：

for epoch in range(start_epoch, nEpochs + 1):
        train(training_data_loader, optimizer, model, criterion, epoch)
        save_checkpoint(model, epoch)

def save_checkpoint(model, epoch):
    model_out_path = "checkpoint/" + "model_epoch_{}.pth".format(epoch)
    state = {"epoch": epoch ,"model": model}
    if not os.path.exists("checkpoint/"):
        os.makedirs("checkpoint/")

    torch.save(state, model_out_path)

    print("Checkpoint saved to {}".format(model_out_path))

上面的代码中start_epoch是开始保存模型的epoch,nEpochs是总共训练的次数。train（）里面的参数，是训练的过程：一些训练数据，优化器，模型和训练标准。