如何在一个模型中训练多个损失，并在 pytorch 中选择冻结网络的某些部分？

Posted 2023-03-13

【中文标题】如何在一个模型中训练多个损失，并在 pytorch 中选择冻结网络的某些部分？

【英文标题】：How to train multiple losses in one model with freezing some part of the network optionally in pytorch?

【发布时间】：2020-10-09 11:49:20

【问题描述】：

我是 pytorch 的新手。 我正在实现一个具有两个分类器（黄色和紫色）​​的网络，如figure 所示。 问题是我想在网络训练黄色分类器时冻结红色部分，在网络训练紫色分类器时解冻红色部分。

我想实现的简短代码如下

# x is input and y_yellow, y_purple are labels of yellow and purple classifiers respectively.
criterion = CrossEntropyLoss()
opt = SGD()
model = my_model()

opt.zero_grad()
yellow_out, purple_out = model(x)

# freeze red part requires_grad = False

yellow_loss = criterion(yellow_out, y_yellow)
yellow_loss.backward()
opt.step()

opt.zero_grad()
# unfreeze red part requires_grad = True

purple_loss = criterion(purple_out, y_purple)
purple_loss.backward()
opt.step

请告诉我实现这个想法的确切方法。

代码顺序对吗？

我是否正确使用了 zero_grad？

我错过了什么吗？

我需要使用任何可选参数吗？

【参考方案1】：

最简单的策略是不将红色模块的参数传递给优化器opt。或者，您可以将红色模块的参数设置为requires_grad 为False。

它看起来像：

for param in red_module.parameters():
    param.requires_grad = False

让我知道这是否适合你。