pytorch 更新部分参数（冻结参数）注意事项

最新推荐文章于 2025-05-04 11:57:52 发布

Answerlzd

最新推荐文章于 2025-05-04 11:57:52 发布

阅读量2.5w

点赞数 49

CC 4.0 BY-SA版权

分类专栏：深度学习入门文章标签： python 神经网络深度学习

本文链接：https://blog.csdn.net/Answer3664/article/details/108493753

实验的pytorch版本1.2.0

在训练过程中可能需要固定一部分模型的参数，只更新另一部分参数。有两种思路实现这个目标，一个是设置不要更新参数的网络层为false，另一个就是在定义优化器时只传入要更新的参数。当然最优的做法是，优化器中只传入requires_grad=True的参数，这样占用的内存会更小一点，效率也会更高。

一、设置参数为false

import torch
import torch.nn as nn
import torch.optim as optim

# 定义一个简单的网络
class net(nn.Module):
    def __init__(self, num_class=10):
        super(net, self).__init__()
        self.fc1 = nn.Linear(8, 4)
        self.fc2 = nn.Linear(4, num_class)
    def forward(self, x):
        return self.fc2(self.fc1(x))


model = net()

# 冻结fc1层的参数
for name, param in model.named_parameters():
    if "fc1" in name:
        param.requires_grad = False


loss_fn = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=1e-2)  # 传入的是所有的参数
print("model.fc1.weight", model.fc1.weight)
print("model.fc2.weight", model.fc2.weight)

for epoch in range(10):
    x = torch.randn((3, 8))
    label = torch.randint(0,10,[3]).long()
    output = model(x)

    loss = loss_fn(output, label)