pytorch 加载部分预训练网络预训练网络比正在训练的网络大完整直接能跑案例_Python

当然说是直接能跑，首先你要有torch

第一个文件，生成存储网络参数的pth文件

网络是预测一个sin函数，全连接，1 20 20 1

# -*- coding: utf-8 -*-
"""
Created on Mon Jul 27 16:47:37 2020
学习PyTorch中使用预训练的模型初始化网络的一部分参数(增减网络层，修改某层参数等) 固定参数
https://discuss.pytorch.org/t/how-to-load-part-of-pre-trained-model/1113/3?u=leyang

这个和test_load_part_of_pretrained_model2 是一个联动
这个文件产生pth
test_load_part_of_pretrained_model2 使用pth
@author: user
"""

import torch
import torch.nn as nn
from torch.utils.data import Dataset,DataLoader
import matplotlib.pyplot as plt
import os
import numpy as np
import random

os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def seed_torch(seed=0):
    random.seed(seed)
    np.random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

seed_torch(2)

#数据处理的类
class MyDataset(Dataset):
    # 创建自己的类：MyDataset,这个类继承的torch.utils.data.Dataset
    #自定义的dataset类需要继承Dataset. 并且重载
    def __init__(self,data,label):
        self.fea=data  
        self.label=label
    
    def __len__(self):
        #返回数据集长度，注意和loader的长度区分开
        return len(self.label)
    
    def __getitem__(self,idx):
        #按照索引读取每个元素的具体内容
        fea=self.fea[idx]
        label=self.label[idx]
        """Convert ndarrays to Tensors."""
        return {'fea':fea,
                'label':label
                }
        #返回的是字典

class NeuralNet(nn.Module):
    def __init__(self,input_size,hidden_size,num_classes):
        super(NeuralNet,self).__init__()
        self.fc1=nn.Linear(input_size,hidden_size)
        self.relu=nn.ReLU()
        self.fc3=nn.Linear(hidden_size,hidden_size)
        self.fc2=nn.Linear(hidden_size,num_classes)
        
        
    def forward(self,x):
        out=self.fc1(x)
        out=self.relu(out)
        out=self.fc3(out)
        out=self.relu(out)
        out=self.fc2(out)
        return out
    


#%%参数
Batch_size=512
input_size=1
hidden_size=20
num_classes=1
learning_rate=0.001
num_epochs=20
#%%数据
y_train=torch.rand(10000,1)
X_train=torch.sin(y_train)
y_test=torch.rand(1000,1)
X_test=torch.sin(y_test)

#创立数据集
train_dataset=MyDataset(X_train,y_train)
test_dataset=MyDataset(X_test,y_test)
#shuffle 打乱顺序 默认为False
train_loader=DataLoader(train_dataset,batch_size=Batch_size,shuffle=True)
test_loader=DataLoader(test_dataset,batch_size=1000)


model=NeuralNet(input_size,hidden_size,num_classes)
criterion=nn.MSELoss()
optimizer=torch.optim.Adam(model.parameters(),lr=learning_rate)
model=model.to(device)
criterion=criterion.to(device)

total_step=len(train_loader)
for epoch in range(num_epochs):
    for i,data in enumerate(train_loader):
        inputs,labels=data['fea'],data['label']
        inputs=inputs.to(device)
        labels=labels.to(device)
        
        outputs=model(inputs)
        loss=criterion(outputs,labels)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if (i+1)%10==0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 
                   .format(epoch+1, num_epochs, i+1, total_step, loss.item()))

loss_func = torch.nn.MSELoss() 
with torch.no_grad():
    correct=0
    total=0
    for i,data in enumerate(test_loader):
        inputs,labels=data['fea'],data['label']
        inputs=inputs.to(device)
        labels=labels.to(device)
        outputs=model(inputs)
        MSE_ = loss_func(outputs, labels)
        MSE=torch.mean((outputs-labels)**2)
        MSE_compare=torch.mean((inputs-labels)**2)
        print(MSE_.item(),MSE.item(),MSE_compare.item())
        plt.cla()
        outputs=outputs.cpu()
        labels=labels.cpu()
        plt.plot([i for i in range(100)],outputs[:100],labels[:100])

state = {'model':model.state_dict(), 'optimizer':optimizer.state_dict(),\
              'epoch':epoch}

torch.save(state, 'test_save.pth')

第二个文件是加载预训练的网络

全连接，网络是1 20 1

# -*- coding: utf-8 -*-
"""
Created on Tue Jul 28 10:20:35 2020
学习PyTorch中使用预训练的模型初始化网络的一部分参数(增减网络层，修改某层参数等) 固定参数
https://discuss.pytorch.org/t/how-to-load-part-of-pre-trained-model/1113/3?u=leyang
test_load_part_of_pretrained_model2 使用pth (使用部分网络)

@author: user
"""



import torch
import torch.nn as nn
from torch.utils.data import Dataset,DataLoader
import matplotlib.pyplot as plt
import os
import numpy as np
import random
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def seed_torch(seed=0):
    random.seed(seed)
    np.random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

seed_torch(2)

#数据处理的类
class MyDataset(Dataset):
    # 创建自己的类：MyDataset,这个类继承的torch.utils.data.Dataset
    #自定义的dataset类需要继承Dataset. 并且重载
    def __init__(self,data,label):
        self.fea=data  
        self.label=label
    
    def __len__(self):
        #返回数据集长度，注意和loader的长度区分开
        return len(self.label)
    
    def __getitem__(self,idx):
        #按照索引读取每个元素的具体内容
        fea=self.fea[idx]
        label=self.label[idx]
        """Convert ndarrays to Tensors."""
        return {'fea':fea,
                'label':label
                }

class NeuralNet(nn.Module):
    def __init__(self,input_size,hidden_size,num_classes):
        super(NeuralNet,self).__init__()
        self.fc1=nn.Linear(input_size,hidden_size)
        self.fc1.bias=nn.Parameter(h_zeros)
        self.relu=nn.ReLU()
        self.fc2=nn.Linear(hidden_size,num_classes)
        
    def forward(self,x):
        out=self.fc1(x)
        out=self.relu(out)
        out=self.fc2(out)
        return out

#%%参数
Batch_size=512
input_size=1
hidden_size=20
num_classes=1
learning_rate=0.001
num_epochs=50
h_zeros=torch.zeros(hidden_size)
#%%数据
y_train=torch.rand(10000,1)
X_train=torch.sin(y_train)
y_test=torch.rand(1000,1)
X_test=torch.sin(y_test)

#创立数据集
train_dataset=MyDataset(X_train,y_train)
test_dataset=MyDataset(X_test,y_test)
#shuffle 打乱顺序 默认为False
train_loader=DataLoader(train_dataset,batch_size=Batch_size,shuffle=True)
test_loader=DataLoader(test_dataset,batch_size=1000)


model=NeuralNet(input_size,hidden_size,num_classes)
model_dict = model.state_dict()
criterion=nn.MSELoss()
optimizer=torch.optim.Adam(model.parameters(),lr=learning_rate)
model=model.to(device)
criterion=criterion.to(device)

#-----核心部分--------------------------------------------------------------
pretrained_dict = torch.load('test_save.pth')
pretrained_model_dict=pretrained_dict['model']
pretrained_model_dict = {k: v for k, v in pretrained_model_dict.items() if k in model_dict}
model_dict.update(pretrained_model_dict)
model.load_state_dict(model_dict)

start_epoch = pretrained_dict['epoch'] + 1
#-----核心部分---------------------------------------------------------------


total_step=len(train_loader)
for epoch in range(start_epoch,num_epochs):
    for i,data in enumerate(train_loader):
        inputs,labels=data['fea'],data['label']
        inputs=inputs.to(device)
        labels=labels.to(device)
        
        outputs=model(inputs)
        loss=criterion(outputs,labels)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if (i+1)%10==0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 
                   .format(epoch+1, num_epochs, i+1, total_step, loss.item()))

loss_func = torch.nn.MSELoss() 
with torch.no_grad():
    correct=0
    total=0
    for i,data in enumerate(test_loader):
        inputs,labels=data['fea'],data['label']
        inputs=inputs.to(device)
        labels=labels.to(device)
        outputs=model(inputs)
        MSE_ = loss_func(outputs, labels)
        MSE=torch.mean((outputs-labels)**2)
        MSE_compare=torch.mean((inputs-labels)**2)
        print(MSE_.item(),MSE.item(),MSE_compare.item())
        plt.cla()
        outputs=outputs.cpu()
        labels=labels.cpu()
        plt.plot([i for i in range(100)],outputs[:100],labels[:100])

注意事项：加载路径是绝对路径（不想暴露我的文件夹所以没贴）

pretrained_dict是个字典

pretrained_model_dict存有网络结构和网络参数，也是个字典，可以使用pretrained_model_dict.keys()查看层数的名字

使用pretrained_model_dict['fc1.weight']查看参数

本文地址：https://blog.csdn.net/Stephanie2014/article/details/107631590