首页 > 脚本专栏 > python > pytorch模型的定义、修改、读取、断点续训

pytorch模型的定义、修改、读取、断点续训深入解析

2024-03-09 09:25:14 作者：nuocheng

模型定义是深度学习中重要的一环,PyTorch提供了强大而灵活的工具和函数,使我们能够轻松定义各种类型的深度学习模型,通过深入理解模型定义的原理和应用,我们能够更好地理解和设计自己的模型,从而提升深度学习任务的性能和效果

在机器学习和深度学习领域，PyTorch是一种广泛使用的开源深度学习框架。它提供了丰富的工具和函数，方便用户定义、训练和部署各种深度学习模型。本篇博客将详细介绍PyTorch中模型定义的方式，并结合原理和代码示例进行讲解，旨在帮助读者深入理解PyTorch的模型定义过程。

前言

模型定义的基本原理

在PyTorch中，模型定义是通过定义一个继承自torch.nn.Module类的Python类来实现的。torch.nn.Module是PyTorch中模型定义的基础，它提供了一组丰富的工具和函数，用于定义和操作神经网络模型。

模型定义的基本原理如下：

创建一个继承自torch.nn.Module的子类，这个子类将成为我们定义的模型。
在子类的构造函数中，首先调用super().__init__()来初始化父类torch.nn.Module，然后在构造函数中定义模型的各个层和模块。
在子类中实现forward方法，该方法定义了模型的前向传播过程，即定义了输入数据如何经过各个层进行计算得到输出。
可选地，在子类中实现__str__方法，用于打印模型的结构信息。
接下来，我们将通过一个简单的神经网络模型的定义和代码示例来进一步解释以上原理。

模型参数和层的概念

在深入了解模型定义之前，让我们先来了解一些基本概念：模型参数和层。

模型参数

模型参数是模型内部可学习的参数，它们会在训练过程中自动更新以优化模型的性能。常见的模型参数包括权重（weights）和偏置（biases）。权重是连接不同神经元的连接强度，而偏置是每个神经元的激活阈值。

层

在PyTorch中，层是模型中的构建块，它们接受输入数据并将其转换为输出数据。层通常包含一些可学习的参数，例如全连接层中的权重和偏置。常见的层类型包括全连接层、卷积层、池化层等。

pytorch的模型定义

pytorch有3种模型定义方式，三种方式都是基于nn.Module建立的，Module是所有网络的基础。

Sequential
ModuleList
ModuleDict

1) Sequential

该方法与tf2很相似，使用也很简单

以数字作为层的名称

import torch
import torch.nn as nn
model = nn.Sequential(
	nn.Linear(56,512),
    nn.ReLU(),
    nn.Linear(512,8)
)
print(model)

对层的名称进行自定义

import collections
import torch
import torch.nn as nn

model = nn.Sequential(collections.OrderedDict([
    ('layer_1',nn.Linear(56,512)),
    ('layer_2',nn.ReLU()),
    ('layer_3',nn.Linear(512,8)),
]))

2) ModuleList

ModuleList接收一个子模块的列表作为输入（一个列表中不同的module，并把参数注册到网络中），可以使用append，extend进行操作。

class ModuleListExample(nn.module):
    def __init__(self):
        super().__init__()
        self.model = nn.ModuleList([nn.Linear(56,512),nn.ReLU(),nn.Linear(512,8)])
    def forward(self,x):
        for layers in self.model:
            x=f(x)
        return x

nn.Sequential内部实现了forward函数，因此可以不用写forward函数，而nn.ModuleList则没有实现内部forward函数，需要人工实现。
nn.Sequential初始调用时就决定了里面Module的调用顺序，因此需要保证上一个Module的输出大小能符合下一个Module输入大小的要求。nn.ModuleList只是简单的将Module打包并注册到网络中，需要在人工实现的forward里面去决定调用顺序。

3)ModuleDict

该方法与ModuleList类似，只是能加入网络层的名称

class ModuleDictExample(nn.module):
    def __init__(self):
        super().__init__()
        self.choices = nn.ModuleDict([
            'conv': nn.Conv2d(10,10,3),
            'pool': nn.MaxPool2d(3)
        ])
        self.activations = nn.ModuleDict([
            ['lrelu',nn.LeakyReLU()],
            ['prelu',nn.PReLU()]
        ])
    def forward(self,x,choice,act):
        x = self.choices[choice](x)
        x = self.activations[act](x)
        return x

模型的修改

1）修改模型

这里以resnet50作为修改模型，修改全连接层的输出大小

import torch.nn as nn
import torchvision.models as models
from collections import OrderedDict

resnet = models.resnet50()

classifier_ten = nn.Sequential(OrderedDict([
      ('layer_1',nn.Linear(56,512)),
      ('layer_2',nn.ReLU()),
      ('layer_3',nn.Linear(512,8)),
]))
net.fc = classifier_ten

2)添加额外输入

将原模型添加输入位置前的部分作为一个整体，同时在forward中定义好原模型不变的部分，添加输入和后续层之间的链接关系；

class myModel(nn.Module):
    def __init__(self,net):
        super().__init__()
        self.net = net
        self.relu = nn.Relu()
        self.dropout = nn.Dropout(0.5)
    def forward(self,x,add_variable):
        x = self.net(x)
        x = torch.cat((self.drop(self.relu(x)),add_variable.unsqueeze(1)))
        x = self.fc_add(x)
        x = self.output(x)
        return x

net = models.resnet50()
print('添加额外输入之前：\n', net)
model = Model(net)
print('添加额外输入之后：\n', model)

3)额外输出

import torch.nn as nn
import  torch
import torchvision.models as models
from collections import OrderedDict

# 添加额外输出
class Model(nn.Module):
    def __init__(self, net):
        super(Model, self).__init__()
        self.net = net
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc1 = nn.Linear(1000, 10, bias=True)
        self.output = nn.Softmax(dim=1)

    def forward(self, x):
        x1000 = self.net(x)
        x10 = self.dropout(self.relu(x1000))
        x10 = self.fc1(x10)
        x10 = self.output(x10)
        # 输出倒数第二层x1000,和最后一层x10
        return x10, x1000

net = models.resnet50()
print('添加额外输出之前：\n', net)
model = Model(net)
print('添加额外输出之后：\n', model)

模型的保存与读取

保存的格式：pkl，pt，pth

多卡模型存储：torch.nn.DataParallel(model).cuda()

1）模型结构与模型参数的保存

保存

save_dir = './models/resnet50.pkl'
model = models.resnet50(pretrained=True)
torch.save(model, save_dir)

加载

loaded_model = toch.load(save_dir)
loaded_model.eval（)

2)只保存权重

保存

save_dir = './models/resnet50_state_dict.pkl'
torch.save(model.state_dict(),save_dir)

model.state_dict()返回的事一个orderdict，存储了网络结构的名字和对应的参数

加载

loaded_dict = torch.load(save_dir) #加载参数
loaded_model = models.resnet50()   #加载模型
loaded_model.state_dict = loaded_dict #还原模型中的变量值
loaded_model.eval（)

loaded_model.state_dict = loaded_dict

可以使用下面的代码进行替换

loaded_model.load_statie_dict(loaded_dict,strict=True)

load_state_dict() 中还有一个关键的参数 strict, 当strict=True,要求预训练权重层数的键值与新构建的模型中的权重层数名称完全吻合；如果新构建的模型在层数上进行了部分微调，则上述代码就会报错：key对应不上, 此时，采用strict=False 就能够解决这个问题

断点续训

模型训练过程中的保存

checkpoint ={
    'net':model.state_dict(),
    'optimzer':optimizer.state_dict(),
    'epoch':epoch
}
torch.save(checkpoint, './models/checkpoint/ckpt_weight_12.pth')

断点回复

1）首先加载最近的权重信息到模型中

2）配置epoch

3）让模型开始训练

path_checkpoint = "./models/checkpoint/ckpt_weight_12.pth"
checkpoint = torch.load(path_checkpoint)
model.load_state_dict(checkpoint['net'])
optimizer.load_state_dict(checkpoint['optimizer'])
start_epoch = checkpoint['epoch']

for epoch in range(start_epoch+1,100):
    for step,(b_img,b_label) in enumerate(train_loader):
        train_output = model(b_img)
        loss = loss_func(train_output,b_label)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

学习率自动下降

TensorFlow提供了一种更加灵活的学习率设置方法-指数衰减法，使用tf.train.exponential_decay实现。指数衰减法的核心思想是，先使用较大的学习率来快速得到一个比较优的解，然后随着迭代的继续逐步减小学习率，使得模型更加稳定。tf.train.exponential_decay函数可以用以下代码实现的功能来展示：

decayed_learning_rate = learning_rate*decay_rate^(global_step/decay_steps)
#decayed_learning_rate为每一轮优化时使用的学习率，learning_rate为事先设定的初始学习率，decay_rate为衰减系数，global_step为迭代次数，decay_steps为衰减速度（即迭代多少次进行衰减）

可见使用此函数，随着迭代次数的增加，学习率逐步降低。

tf.train.exponential_decay可以通过设置参数staircase选择不同的衰减方式，其默认值为False，既每一次迭代都进行学习率的优化，

global_step = tf.Variable(0)
learning_rate = tf.train.exponential_decay(0.1,global_step,100,0.96,staircase=True)
#0.1表示初始学习率，global_step表示迭代次数，100表示衰减速度，0.96表示衰减率
#使用指数衰减的学习率，在minimize函数中传入global_step，它将自动更新，learning_rate也随即被更新
learning_step = tf.train.GradientDescentOptimizer(learning_rate).minimize(loss, global_step=global_step)
#神经网络反向传播算法，使用梯度下降算法GradientDescentOptimizer来优化权重值，learning_rate为学习率，minimize中参数loss是损失函数，global_step表明了当前迭代次数(会被自动更新)

pytorch中进行使用

import torch
import torch.nn as nn
from torch.optim.lr_scheduler import StepLR
import itertools
initial_lr = 0.1

class model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3)

    def forward(self, x):
        pass

net_1 = model()
optimizer_1 = torch.optim.Adam(net_1.parameters(),lr=initial_lr)
#对优化方法进行优化
scheduler_1 = StepLR(optimizer_1, step_sie=3, gamma=0.1)
for epoch in range(1, 11):

    optimizer_1.zero_grad()
    optimizer_1.step()
    print("第%d个epoch的学习率：%f" % (epoch, optimizer_1.param_groups[0]['lr']))
    scheduler_1.step()

import torch
import torch.nn as nn
from torch.optim.lr_scheduler import LambdaLR
class model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3)

    def forward(self, x):
        pass

net_1 = model()

optimizer_1 = torch.optim.Adam(net_1.parameters(), lr = initial_lr)
scheduler_1 = LambdaLR(optimizer_1, lr_lambda=lambda epoch: 1/(epoch+1))

print("初始化的学习率：", optimizer_1.defaults['lr'])

for epoch in range(1, 11):
    # train

    optimizer_1.zero_grad()
    optimizer_1.step()
    print("第%d个epoch的学习率：%f" % (epoch, optimizer_1.param_groups[0]['lr']))
    scheduler_1.step()

多步长SGD继续训练

在简单的任务中，我们使用固定步长（也就是学习率LR）进行训练，但是如果学习率lr设置的过小的话，则会导致很难收敛，如果学习率很大的时候，就会导致在最小值附近，总会错过最小值，loss产生震荡，无法收敛。所以这要求我们要对于不同的训练阶段使用不同的学习率，一方面可以加快训练的过程，另一方面可以加快网络收敛。

所以我们在保存网络中的训练的参数的过程中，还需要保存lr_scheduler的state_dict，然后断点继续训练的时候恢复

optimizer = torch.optim.SGD(model.parameters(),lr=0.1)
#这里我设置了不同的epoch对应不同的学习率衰减，在10->20->30，学习率依次衰减为原来的0.1，即一个数量级
lr_schedule = torch.optim.lr_scheduler.MultiStepLR(optimizer,milestones=[10,20,30,40,50],gamma=0.1)


for epoch in range(start_epoch+1,80):
    optimizer.zero_grad()
    optimizer.step()
    lr_schedule.step()

    if epoch %10 ==0:
        print('epoch:',epoch)
        print('learning rate:',optimizer.state_dict()['param_groups'][0]['lr'])

我们在保存的时候，也需要对lr_scheduler的state_dict进行保存，断点继续训练的时候也需要恢复lr_scheduler

if RESUME:
    path_checkpoint = "./model_parameter/test/ckpt_best_50.pth"  # 断点路径
    checkpoint = torch.load(path_checkpoint)  # 加载断点

    model.load_state_dict(checkpoint['net'])  # 加载模型可学习参数

    optimizer.load_state_dict(checkpoint['optimizer'])  # 加载优化器参数
    start_epoch = checkpoint['epoch']  # 设置开始的epoch
    lr_schedule.load_state_dict(checkpoint['lr_schedule'])#加载lr_scheduler
for epoch in range(start_epoch+1,80):

    optimizer.zero_grad()

    optimizer.step()
    lr_schedule.step()


    if epoch %10 ==0:
        print('epoch:',epoch)
        print('learning rate:',optimizer.state_dict()['param_groups'][0]['lr'])
        checkpoint = {
            "net": model.state_dict(),
            'optimizer': optimizer.state_dict(),
            "epoch": epoch,
            'lr_schedule': lr_schedule.state_dict()
        }
        if not os.path.isdir("./model_parameter/test"):
            os.mkdir("./model_parameter/test")
        torch.save(checkpoint, './model_parameter/test/ckpt_best_%s.pth' % (str(epoch)))

optimizer = torch.optim.SGD(model.parameters(),lr=0.1)
lr_schedule = torch.optim.lr_scheduler.MultiStepLR(optimizer,milestones=[10,20,30,40,50],gamma=0.1)
start_epoch = 9
# print(schedule)


if RESUME:
    path_checkpoint = "./model_parameter/test/ckpt_best_50.pth"  # 断点路径
    checkpoint = torch.load(path_checkpoint)  # 加载断点

    model.load_state_dict(checkpoint['net'])  # 加载模型可学习参数

    optimizer.load_state_dict(checkpoint['optimizer'])  # 加载优化器参数
    start_epoch = checkpoint['epoch']  # 设置开始的epoch
    lr_schedule.load_state_dict(checkpoint['lr_schedule'])

for epoch in range(start_epoch+1,80):

    optimizer.zero_grad()

    optimizer.step()
    lr_schedule.step()


    if epoch %10 ==0:
        print('epoch:',epoch)
        print('learning rate:',optimizer.state_dict()['param_groups'][0]['lr'])
        checkpoint = {
            "net": model.state_dict(),
            'optimizer': optimizer.state_dict(),
            "epoch": epoch,
            'lr_schedule': lr_schedule.state_dict()
        }
        if not os.path.isdir("./model_parameter/test"):
            os.mkdir("./model_parameter/test")
        torch.save(checkpoint, './model_parameter/test/ckpt_best_%s.pth' % (str(epoch)))

总结

模型定义是深度学习中重要的一环，PyTorch提供了强大而灵活的工具和函数，使我们能够轻松定义各种类型的深度学习模型。通过深入理解模型定义的原理和应用，我们能够更好地理解和设计自己的模型，从而提升深度学习任务的性能和效果。

到此这篇关于pytorch模型的定义、修改、读取、断点续训深入解析的文章就介绍到这了,更多相关pytorch模型的定义、修改、读取、断点续训内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家！