当前位置：首页 > article >正文

如何固定权重，对某些层得学习率改为0？

article 2024/9/24 4:20:02

PyTorch 源码解读之 torch.autograd：梯度计算详解
在训练深度神经网络时，有时需要固定某些层或参数不进行更新。可以使用PyTorch提供的nn.Module中的parameters()方法来获得所有需要学习的参数，并使用torch.optim.SGD()等优化器的param_groups参数来控制不同层的学习率。通过将学习率设为0，就可以固定某些层或参数的权重。

例如，在以下示例中，假设我们要固定某个网络的前3个卷积层的权重，不更新这些层的权重：

import torch.optim as optim

model = MyModel()  # 这里定义自己的模型
optimizer = optim.SGD([{'params': model.fc.weight, 'lr': 0.01},
                       {'params': model.fc.bias, 'lr': 0.02},
                       {'params': model.conv1.parameters()},
                       {'params': model.conv2.parameters(), 'lr': 0}], lr=0.1)

for epoch in range(num_epochs):
    # 训练代码

在上面的代码中，我们将第1个全连接层的权重设为0.01的学习率，第1个全连接层的偏置设为0.02的学习率。我们还将第1个卷积层和第2个卷积层的权重都设置为0的学习率。这将会固定前3个卷积层的权重，同时调整全连接层的学习率。在训练过程中，只会更新全连接层的权重，而不会更新前3个卷积层。

另一种设置不同层学习率的方法是使用参数组。在PyTorch中，优化器的参数组用于对不同的参数设置不同的学习率和不同的权重衰减。可以使用optim.SGD()等优化器的param_groups参数，将不同的参数划分到不同的组中，进而可以针对不同的层分别设置学习率。

例如，在以下示例中，假设我们要将某个网络的前3个卷积层的权重固定，并且不进行权重衰减，而对全连接层的权重进行学习率为0.01，权重衰减为0.001的设置：

import torch.optim as optim

model = MyModel()  # 这里定义自己的模型
params = [{'params': model.conv1.parameters()},
          {'params': model.conv2.parameters()},
          {'params': model.conv3.parameters()},
          {'params': model.fc.parameters(), 'lr': 0.01, 'weight_decay': 0.001}]

optimizer = optim.SGD(params, lr=0.1)

for epoch in range(num_epochs):
    # 训练代码

在这个例子中，我们首先将前3个卷积层的参数分为一组，使用model.conv1.parameters()等方法来获取这些参数。然后，我们将全连接层的参数分为另外一组，并通过指定lr和weight_decay参数来设置不同的学习率和权重衰减。最后，我们将这些组合成一个参数列表，并传递给optim.SGD()优化器。在训练过程中，前三个卷积层的权重会被固定，并且不会进行权重衰减，而全连接层的权重则会按照我们设置的学习率和权重衰减进行更新。

除了在代码中手动设置不同层的学习率之外，还可以使用一些自适应学习率调整的优化器，例如Adam、Adagrad、Adadelta等。这些优化器可以自动调整每个参数的学习率，使得不同参数能够以适当的学习率进行更新，具有较好的性能表现。

以Adam为例，它的学习率动态调整的方法是根据梯度的平方均值、梯度的一阶矩估计和二阶矩估计来计算每个参数的学习率。在PyTorch中，可以使用torch.optim.Adam()来创建Adam优化器。

以下示例展示了如何使用Adam优化器进行训练，并自动调整不同参数的学习率：

import torch.optim as optim

model = MyModel()  # 这里定义自己的模型
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(num_epochs):
    # 训练代码
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

在这个例子中，我们使用optim.Adam()创建了一个Adam优化器，将模型的所有参数传递给它。在每个epoch中，我们先使用optimizer.zero_grad()将梯度清零，然后进行正向传播和反向传播计算损失和梯度，最后使用optimizer.step()更新所有参数，并自动调整每个参数的学习率。

需要注意的是，使用自适应学习率的优化器，不再需要手动设置不同层的学习率，优化器会根据每个参数的情况自动进行调整，具有很好的鲁棒性和泛化性能。

PyTorch

在PyTorch中，可以使用requires_grad属性来控制模型中哪些参数需要计算梯度，从而控制哪些层的权重参与训练，哪些层的权重不参与训练。

具体来说，如果需要固定某些层的权重，可以将其对应的参数的requires_grad属性设置为False，即不需要计算梯度。这样，在模型的反向传播过程中，这些层的权重就不会更新。

以下是一个示例代码，演示如何固定模型的前两层的权重不参与训练：

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layer1 = nn.Linear(10, 20)
        self.layer2 = nn.Linear(20, 30)
        self.layer3 = nn.Linear(30, 40)
        
    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        return x

model = MyModel()

# 固定前两层的权重不参与训练
for param in model.layer1.parameters():
    param.requires_grad = False
for param in model.layer2.parameters():
    param.requires_grad = False

# 打印模型参数的requires_grad属性
for name, param in model.named_parameters():
    print(name, param.requires_grad)