pytorch不加载fc_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']查看参数

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

如您对本文有疑问或者有任何想说的，请点击进行留言回复，万千网友为您解惑！