PyTorch：利用预训练好的ResNet-152网络提取图片特征

这篇博文相当于是对上一篇博文Pytorch：利用预训练好的VGG16网络提取图片特征 的补充，本文中提到的提取方式与前文中的不同。

另外，因为TorchVision提供的训练好了的ResNet效果不好，所以本文中将会使用由ruotianluo提供的从Caffe转换过来的ResNet模型（具体可以看这个repo，如果好奇怎么转换的话）。

代码

以下代码节选自pytorch-vqa的preprocess-images.py，作者是Cyanogenoid。

import h5py
from torch.autograd import Variable
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torch.utils.data
import torchvision.models as models
from tqdm import tqdm

import config
import data
import utils
from resnet import resnet as caffe_resnet

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.model = caffe_resnet.resnet152(pretrained=True)

        def save_output(module, input, output):
            self.buffer = output
        self.model.layer4.register_forward_hook(save_output)

    def forward(self, x):
        self.model(x)
        return self.buffer

上面就是一个例子，重点就在于self.model.layer4.register_forward_hook(save_output)（代码中layer4就是$2048,14,14$大小的feature所在的层，即最后一层卷积层），这其实是官方推荐的提取方式，可以想象成用了个hook钩子来把对应层的feature钓了出来。用法就是这样：

net = Net().cuda()
net.eval()

feat = net(img)	# img是(1, 3, 448, 448)大小的，1是batch size，这里只有一张图片，3是RGB通道
				# 如果只有1要注意预处理，两个448就是图片大小了，不符合的话要记得resize

另一种提取方式

这里提到另一种方法主要是因为我在提取最后一层pooling层的时候没办法用上面这个方法，所以从另外一个地方找到了这个方法。

首先，将模型的最后一层全连接层删除掉，然后再让最后一层变成x=x的直接映射的层，最后直接喂数据，得到feature。

# 这里省略掉一堆import
import torchvision.models as models

from resnet import resnet as caffe_resnet

# 省略掉读取图片和预处理的步骤，下面的img就是已经经过预处理之后的图片

model = caffe_resnet.resnet152(pretrained=True)
del model.fc
model.fc=lambda x:x
model = model.cuda()

feat = model(img)

补充

最后补充一点，一定要注意预处理方式，本文中用到的模型需要的预处理方式是：

def get_transform(target_size, central_fraction=1.0):
    return transforms.Compose([
        transforms.Scale(int(target_size / central_fraction)),
        transforms.CenterCrop(target_size),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225]),
])

transform = get_transform(448 // 32, 0.85)
img = transform (img)

参考资料

pytorch-vqa：很有用的repo
Extract a feature vector for any image with PyTorch：更加详细的介绍，但是不是ResNet，不过也通用