`optimizer.step()` 如何执行就地操作？

Question

这是一个导致就地操作错误的简单示例。

import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import OrderedDict
from torch import optim

torch.autograd.set_detect_anomaly(True)

class Loss(nn.Module):
    def __init__(self):
        super(Loss, self).__init__()
    def forward(self, x, target):
        return x[0,0,0,0]

def block(in_channels, features, name):
    return nn.Conv2d(in_channels=in_channels,
                     out_channels=features,
                     kernel_size=3,
                     padding=1,
                     bias=False)
class SharedNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.shared_layer = block(in_channels=3, features=1, name="wow")

    def forward(self, x):
        x = self.shared_layer(x)
        return x


class Network1(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = block(in_channels=1, features=1, name="wow-1")
    def forward(self, x):
        return self.conv(x)
    
class Network2(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = block(in_channels=1, features=1, name="wow-2")
    def forward(self, x):
        return torch.sigmoid(self.conv(x))

shared_net = SharedNetwork()
net_1 = Network1()
segmentor = Network2()

optimizer = optim.Adam(list(shared_net.parameters()) + list(segmentor.parameters()), lr=1e-6)
optimizer_conf = optim.Adam(list(shared_net.parameters()), lr=1e-6)
loss_fn = Loss()

# 2. Run a forward pass
fake_data = torch.randint(0,255,(1, 3, 256, 256))/255
target_data_1 = torch.randint(0,255,(1, 3, 256, 256))/255
target_data_2 = torch.randint(0,255,(1, 3, 256, 256))/255


optimizer.zero_grad()
optimizer_conf.zero_grad()

features = shared_net(fake_data)
segmented = segmentor(features)
s_loss = loss_fn(segmented, target_data_2)
s_loss.backward(retain_graph=True)
optimizer.step()


out_1 = net_1(features)
loss = loss_fn(out_1, target_data_1)
loss.backward(retain_graph=False)
optimizer_conf.step()

错误信息：

UserWarning: Error detected in ConvolutionBackward0. No forward pass information available. Enable detect anomaly during forward pass for more information. (Triggered internally at C:\cb\pytorch_1000000000000\work\torch\csrc\autograd\python_anomaly_mode.cpp:97.)
  Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass

one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [1, 3, 3, 3]] is at version 2; expected version 1 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!

我能够通过更改运行优化程序step function 的顺序来解决问题。

optimizer_conf.zero_grad()
optimizer.zero_grad()

features = shared_net(fake_data)
segmented = segmentor(features)
s_loss = loss_fn(segmented, target_data_2)
s_loss.backward(retain_graph=True)


out_1 = net_1(features)
loss = loss_fn(out_1, target_data_1)
loss.backward(retain_graph=False)
optimizer_conf.step()
optimizer.step()

但是，以下问题仍然存在：

1. step方法如何导致卷积中的inplace操作？
2. 为什么将步骤移到文件末尾可以解决此错误？

注意：损失 function 是为了简单起见，使用 dice-loss 也会导致同样的错误！

Answer 1

在回答这个问题之前，我不得不提一下，对一组参数使用多个优化器似乎不是反模式，最好避免。

1. step方法如何导致卷积中的inplace操作？ A: step 方法将梯度添加到权重中，因此它会执行如下操作：

param.weight += param.grad

可以解释为就地操作

2. 为什么将步骤移到文件末尾可以解决此错误？ A：显然，将step方法移到第二个backward方法之后，并没有执行上述操作。 因此，没有就地操作，也没有因它们的存在而引发的错误。

综上所述，一组参数最好只有一个优化器，前面的例子可以这样编码：

import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import OrderedDict
from torch import optim

torch.autograd.set_detect_anomaly(True)

class Loss(nn.Module):
    def __init__(self):
        super(Loss, self).__init__()
    def forward(self, x, target):
        return x[0,0,0,0]
def block(in_channels, features, name):
    return nn.Conv2d(in_channels=in_channels,
                     out_channels=features,
                     kernel_size=(3,3),
                     padding=1,
                     bias=False)
class SharedNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.shared_layer = block(in_channels=3, features=1, name="wow")

    def forward(self, x):
        x = self.shared_layer(x)
        return x


class Network1(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = block(in_channels=1, features=1, name="wow-1")
    def forward(self, x):
        return self.conv(x)
    
class Network2(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = block(in_channels=1, features=1, name="wow-2")
    def forward(self, x):
        return torch.sigmoid(self.conv(x))
    
torch.manual_seed(0)

shared_net = SharedNetwork()
net_1 = Network1()
net_2 = Network2()

shared_optimizer = optim.Adam(list(shared_net.parameters()), lr=1e-6) 
net_1_optimizer = optim.Adam(list(net_1.parameters()), lr=1e-6)
net_2_optimizer = optim.Adam(list(segmentor.parameters()), lr=1e-6)
loss_fn = Loss()

# 2. Run a forward pass
fake_data = torch.randint(0,255,(1, 3, 256, 256))/255
target_data_1 = torch.randint(0,255,(1, 3, 256, 256))/255
target_data_2 = torch.randint(0,255,(1, 3, 256, 256))/255

net_2_optimizer.zero_grad()
features = shared_net(fake_data)
net_2_out = net_2(features)
s_loss = loss_fn(net_2_out, target_data_2)
s_loss.backward(retain_graph=True)
net_2_optimizer.step()

net_1_optimizer.zero_grad()
shared_optimizer.zero_grad()
out_1 = net_1(features)
loss = loss_fn(out_1, target_data_1)
loss.backward(retain_graph=False)
net_1_optimizer.step()
shared_optimizer.step()

注意：如果你想对一组参数应用不同损失的两个不同学习率，你可以根据损失的重要性乘以一个值。 例如，您可以将 loss_1 乘以 0.1，将 loss_1 乘以 0.5。 或者，您可以使用此评论中提到的向后挂钩： backward-hook