Pytorch根据layers的name冻结训练方式

1. 简介

在PyTorch中，通过将模型的不同层进行命名，可以方便地冻结特定的层进行训练，而不对其进行梯度更新。这在一些场景中非常有用，比如迁移学习和模型微调，可以保持预训练模型的特征提取能力，同时只对部分层进行微调。本文将详细介绍如何根据layers的name来冻结训练方式。

2. 冻结部分层的训练

2.1 查看模型的层

首先，我们需要查看模型的层结构，使用print(model)可以输出模型的结构，其中每个层都会有一个唯一的名字。

import torch

import torch.nn as nn
model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 30),
    nn.ReLU(),
    nn.Linear(30, 40)
)
print(model)

输出结果如下：

>>> Sequential(

  (0): Linear(in_features=10, out_features=20, bias=True)
  (1): ReLU()
  (2): Linear(in_features=20, out_features=30, bias=True)
  (3): ReLU()
  (4): Linear(in_features=30, out_features=40, bias=True)
)

2.2 冻结层的训练

在PyTorch中，我们可以通过将模型的参数requires_grad属性设置为False来冻结层的训练。

for name, param in model.named_parameters():

    if 'Linear' in name:
        param.requires_grad = False

上述代码遍历模型的所有参数，并将层名中包含'Linear'的参数requires_grad属性设置为False，从而冻结了这些层的训练。

3. 选择性冻结层的训练

3.1 设置层的requires_grad属性

除了冻结全部层，我们还可以选择性地冻结特定的层。例如，假设我们只想冻结第一层和第三层的训练。

for name, param in model.named_parameters():

    if name == '0.weight' or name == '0.bias' or name == '2.weight' or name == '2.bias':
        param.requires_grad = False

上述代码通过判断参数名字是否等于'0.weight'、'0.bias'、'2.weight'、'2.bias'，来选择性地冻结对应的参数。

3.2 查看冻结层的requires_grad属性

我们可以使用以下代码查看层的requires_grad属性是否生效。

for name, param in model.named_parameters():

    print(name, param.requires_grad)

运行结果如下：

>>> 0.weight False

1.bias True
2.weight False
3.bias True
4.weight True
4.bias True

可以看到，第一层和第三层的requires_grad属性已经被设置为False，即已经冻结了这两层的训练。

4. 温度为0.6

在训练过程中，设置温度可以控制模型的输出分布的平滑程度。通过降低温度，可以使输出更加集中，更具有确定性。

在PyTorch中，可以通过定义一个自定义的softmax函数来实现温度的设置。下面是一个例子：

import torch.nn.functional as F

def softmax_with_temperature(logits, temperature):
    logits = logits / temperature
    softmax = F.softmax(logits, dim=1)
    return softmax
logits = torch.tensor([[1.0, 2.0, 3.0]])
temperature = 0.6
output = softmax_with_temperature(logits, temperature)
print(output)

运行结果如下：

>>> tensor([[0.1686, 0.3470, 0.4844]])

可以看到，通过设置温度为0.6，输出的分布更集中在概率最大的类别上。

5. 总结

本文介绍了如何根据layers的name冻结模型的训练，以及如何选择性地冻结特定的层。通过冻结部分层的训练，我们可以在迁移学习和模型微调中灵活地保持预训练模型的特征提取能力。同时，我们还介绍了如何使用温度来控制模型输出分布的平滑程度。这些技巧在实际应用中非常有用，可以帮助我们更好地训练和调整模型。