PyTorch中如何冻结指定层？DR类仅更新linear3实现方案

给定以下PyTorch模型类：

import torch
import torch.nn as nn
import torch.nn.functional as F

class DR(nn.Module):
    def __init__(self, orginal, latent_dims):
        super(DR, self).__init__()
        self.latent_dims=latent_dims
        self.linear1 = nn.Linear(orginal, 1000)
        self.linear2 = nn.Linear(1000, 2000)
        self.linear3 = nn.Linear(2000, latent_dims)

    def forward(self, x):
        x = F.relu(self.linear1(x))
        x = F.relu(self.linear2(x))
        x = F.relu(self.linear3(x))
        return x

希望训练时linear1和linear2的参数保持初始化状态不更新，仅对linear3的参数进行更新，该如何实现？

解决方案 #

方法一：设置参数的requires_grad属性为False #

通过手动关闭linear1和linear2参数的梯度计算，反向传播时这些参数不会被更新：

# 初始化模型（替换为你的实际输入维度和latent维度）
model = DR(orginal=512, latent_dims=128)

# 冻结linear1和linear2的参数
for param in model.linear1.parameters():
    param.requires_grad = False
for param in model.linear2.parameters():
    param.requires_grad = False

# 定义优化器，仅更新requires_grad=True的参数（即linear3的参数）
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

方法二：优化器仅传入需要更新的参数 #

直接在优化器中指定仅对linear3的参数进行更新，无需修改参数的梯度属性：

# 初始化模型（替换为你的实际输入维度和latent维度）
model = DR(orginal=512, latent_dims=128)

# 仅将linear3的参数传入优化器
optimizer = torch.optim.Adam(model.linear3.parameters(), lr=1e-3)

关键说明 #

• 方法一适合需要后续灵活解冻层的场景，只需将requires_grad重新设为True即可恢复对应层的训练。
• 方法二更简洁直接，避免了遍历参数的操作，性能上更高效。

内容的提问来源于stack exchange，提问作者Nusrat Jahan