当然说是直接能跑,首先你要有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
如您对本文有疑问或者有任何想说的,请点击进行留言回复,万千网友为您解惑!