首页 > 解决方案 > 向层添加前向钩子是否确保使用层的输出计算的损失梯度会自动计算?

问题描述

我有一个模型

class NewModel(nn.Module):
    def __init__(self,output_layer,*args):
        self.output_layer = output_layer
        super().__init__(*args)

        self.output_layer = output_layer
        self.selected_out = None
        #PRETRAINED MODEL
        self.pretrained = models.resnet18(pretrained=True)
    
        #TAKING OUTPUT FROM AN INTERMEDIATE LAYER

        #self._layers = []
        for l in list(self.pretrained._modules.keys()):
            #self._layers.append(l)
            if l == self.output_layer:
                handle = getattr(self.pretrained,l).register_forward_hook(self.hook)
   
    def hook(self,module, input,output):
        self.selected_out = output

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

我有两个目标输出,一个与图像的任何标签相同,第二个与从 获得的输出具有相同的尺寸self.output_layer,称为target_feature

out = model(img)
layerout = model.selected_out

现在,如果我想用目标特征图计算 layerout 的损失,可以像下面写的那样做吗?

loss = criterion(y_true, out) + feature_criterion(layerout, target_feature)

还是我需要添加backward_hooks

在这个 Kaggle 笔记本中

https://www.kaggle.com/sironghuang/understanding-pytorch-hooks

写的是用的loss.backward()时候不能用backward_hooks

引用作者

# backprop once to get the backward hook results
out.backward(torch.tensor([1,1],dtype=torch.float),retain_graph=True)
#! loss.backward(retain_graph=True)  # doesn't work with backward hooks, 
#! since it's not a network layer but an aggregated result from the outputs of last layer vs target 

那么如何根据损失函数计算梯度呢?

标签: pythonpytorch

解决方案


如果我理解正确,您希望从模型中获得两个输出,计算两个损失,然后将它们组合并反向传播。我想你从尝试实现它的方式来自 Tensorflow 和 Keras。在 Pytorch 中,它实际上是相当直接的,你可以很容易地做到这一点,因为它纯粹是功能方面的。

这只是一个例子:

class NewModel(nn.Module):
    def __init__(self, output_layer, *args):
        super(MyModel, self).__init__()
        
        self.pretrained = models.resnet18(pretrained=True)
        self.output_layer = output_layer
        
    def forward(self, x):
        out = self.pretrained(x)
        features = self.output_layer(out)
        return out, features

在推理时,每次调用您将获得两个结果:

>>> m = NewModel(nn.Linear(1000, 10))
>>> x = torch.rand(16, 3, 224, 224)
>>> y_pred, y_feature = m(x)

称你为损失函数:

>>> loss = criterion(y_pred, y_true) + feature_criterion(y_feature, target_feature)

然后,反向传播loss.backward()

所以你的调用不需要钩子,也不需要复杂的渐变.backward


编辑- 如果您想提取中间层输出,请保留钩子,这很好。只需修改forward定义。

def forward(self, x):
    out = self.pretrained(x)
    return out, self.selected_out

例如:

>>> m = NewModel(output_layer='layer1')
>>> x = torch.rand(16, 3, 224, 224)
>>> y_pred, y_feature = m(x)

>>> y_pred.shape, y_feature.shape
(torch.Size([16, 1000]), torch.Size([16, 64, 56, 56]))

另外,我上面所说的关于损失剧照的内容。计算你的损失,然后调用loss.backward()


推荐阅读