记使用Pytorch的hook机制提取特征时踩的一个坑

因为项目需求，需要用DenseNet模型提取图片特征，在使用Pytorch的hook机制提取特征，调试的时候发现提取出来的特征数值上全部大于等于0。
在这里插入图片描述很明显提取出来的特征是经过ReLU的。
现在来看一下笔者是怎么定义hook的：

fmap_block = []
# 注册hook
def forward_hook(module, input, output):
    fmap_block.append(output)

get_feature_model = densenet121(num_classes=2, pretrained=False)
model_dict = torch.load(model_weight_path)
get_feature_model = nn.DataParallel(get_feature_model.cuda())
get_feature_model.module.features.register_forward_hook(forward_hook)

模型定义的时候因项目需求，笔者并没有使用预训练模型。而是自己训练了一个DenseNet121模型，并且使用了DataParallel进行包装。这里有两点需要注意：

大部分的官方模型都会分成两部分，分别是特征层和分类层。

    def forward(self, x):
        features = self.features(x)
        out = F.relu(features, inplace=True)
        out = F.avg_pool2d(out, kernel_size=7, stride=1).view(features.size(0), -1)
        out = self.classifier(out)
        return out

这是DenseNet模型前向传播代码，很明显就是笔者上诉说的那样。所以在使用Pytorch的hook进行提取特征的时候可以很方便的定义成这个样子：

DenseNet类实例.features.register_forward_hook(forward_hook)

眼尖的读者可以发现笔者的代码里并不是这样定义的，多了一个.module（这也算是一个小小的坑）。这是因为笔者使用了DataParallel进行包装模型，使之可以使用多GPU训练，下面来看一下DataParallel的源码：

    def __init__(self, module, device_ids=None, output_device=None, dim=0):
        super(DataParallel, self).__init__()

        if not torch.cuda.is_available():
            self.module = module
            self.device_ids = []

可以看到初始化DataParallel类的时候，将model作为一个参数传给了module，所以得多加一个.module才能定位到我们需要的feature。

看到这里，估计很多人已经发现问题在哪里了，没错，问题出现了前向传播部分，更准确的来说是relu函数。

out = F.relu(features, inplace=True)

inplace表示原地修改张量，所以经过relu层时提前放在列表中的特征张量就会被修改。
两种解决方法：

将inplace置为False，这样就不会原地修改张量了。
修改hook函数

def forward_hook(module, input, output):
    fmap_block.append(output.detach().cpu())

记使用Pytorch的hook机制提取特征时踩的一个坑

Guess you like