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

当然说是直接能跑，首先你要有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']查看参数