pytorch加载已有模型的方式以及使用+加载预训练模型进行部分参数复制

一、加载已有模型直接使用

temp=torch.load("E:\\study-proj\\图像分类：从零到亿\\5.使用更多模型\\model_resnet101.pth") #加载模型，如果只有数值就只会加载模型数据，如果有字典，则会加载模型数据和字典数据
model.load_state_dict(temp)  #返回是否成功

由于模型保存的时候有保存数据和保存数据和字典的方式，所以加载的时候就有两种，利用torch.load，可将不管是数据还是数据和字典都可以加载上，但是如果只是数据，就需要将数据加载到对应的模型上，所以就有如下两种方式：

加载字典（模型对应的字典，其实就是模型变量）以及数据

model=torch.load("pth文件路径")

加载数据

temp=torch.load("E:\\study-proj\\图像分类：从零到亿\\5.使用更多模型\\model_resnet101.pth") #加载模型，如果只有数值就只会加载模型数据，如果有字典，则会加载模型数据和字典数据
model.load_state_dict(temp)  #返回是否成功

完整示例：
代码位置：https://gitee.com/sxh_and_ll/AI-CV/blob/master/proj/%E4%BD%BF%E7%94%A8pytorch%E8%87%AA%E5%B8%A6%E6%A8%A1%E5%9E%8B%E8%AE%AD%E7%BB%83%E4%BB%A5%E5%8F%8A%E6%B5%8B%E8%AF%95/load_model_test.py

'''
    加载模型，进行测试
'''
import time
import torch
from torch import nn
from torch.utils.data import DataLoader
from utils import LoadData

from torchvision.models import alexnet  #最简单的模型
from torchvision.models import vgg11, vgg13, vgg16, vgg19   # VGG系列
from torchvision.models import resnet18, resnet34,resnet50, resnet101, resnet152    # ResNet系列
from torchvision.models import inception_v3     # Inception 系列


def test(dataloader, model):
    size = len(dataloader.dataset)
    # 将模型转为验证模式
    model.eval()
    # 初始化test_loss 和 correct， 用来统计每次的误差
    test_loss, correct = 0, 0
    # 测试时模型参数不用更新，所以no_gard()
    # 非训练， 推理期用到
    with torch.no_grad():
        # 加载数据加载器，得到里面的X（图片数据）和y(真实标签）
        for X, y in dataloader:
            # 将数据转到GPU
            X, y = X.to(device), y.to(device)
            # 将图片传入到模型当中就，得到预测的值pred
            pred = model(X)
            # 计算预测值pred和真实值y的差距
            test_loss += loss_fn(pred, y).item()
            # 统计预测正确的个数
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()#返回相应维度的最大值的索引
    test_loss /= size
    correct /= size
    print(f"correct = {correct}, Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")




if __name__=='__main__':
    batch_size = 8

    # # 给训练集和测试集分别创建一个数据集加载器
    train_data = LoadData("train.txt", True)
    valid_data = LoadData("test.txt", False)

    test_dataloader = DataLoader(dataset=valid_data, num_workers=4, pin_memory=True, batch_size=batch_size)

    # 如果显卡可用，则用显卡进行训练
    device = "cuda" if torch.cuda.is_available() else "cpu"
    # device='cpu'
    print(f"Using {device} device")


    # 加载模型
    temp=torch.load("E:\\study-proj\\图像分类：从零到亿\\5.使用更多模型\\model_resnet101.pth") #加载模型，如果只有数值就只会加载模型数据，如果有字典，则会加载模型数据和字典数据
    model.load_state_dict(temp)
    print(model)
    
    # 定义损失函数，计算相差多少，交叉熵，
    loss_fn = nn.CrossEntropyLoss()

    # 定义优化器，用来训练时候优化模型参数，随机梯度下降法
    optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)  # 初始学习率

    epochs = 1
    for t in range(epochs):
        print(f"Epoch {t+1}\n-------------------------------")
        # print(f"train time: {(time_end-time_start)}")
        time_start=time.time()
        test(test_dataloader, model)
        time_end=time.time()
    print("Done!")

二、加载预训练模型进行参数复制

相关定义

state_dict()获得模型的网络以及参数字典

实现流程

1. 获得相关模型以及预训练的模型参数字典；
2. 循环遍历网络层，一致便复制相关参数，否则就不操作；

#加载预训练模型
resnet = models.resnet50(pretrained=True)
new_state_dict = resnet.state_dict()
dd = net.state_dict()
for k in new_state_dict.keys():
    print(k)
    if k in dd.keys() and not k.startswith('fc'):
        print('yes')
        dd[k] = new_state_dict[k]
net.load_state_dict(dd)