Pytorch：利用预训练好的VGG16网络提取图片特征

注意

目前（2019年一月）因为torchvision提供的VGG网络没有训练完全¹，不建议使用torchvision提供的预训练模型来进行特征提取，建议先使用别的框架（例如TensorFlow或者Caffe之类的框架）提供的预训练过的模型来进行特征提取。

前言

这里的提取图片特征特指从VGG网络的最后一个conv层进行提取。虽然下面代码里面给出的是VGG16作为例子，其实也可以用其他的已经经过训练了的神经网络，包括自己训练的。

相关代码

模型结构

首先说下加载模型，这里用的是torch官方提供的已经训练好的模型，只需要从torchvision模块导入：

import torchvision.models as models

model = models.vgg16(pretrained=True)

上面的pretrained=True是指使用预训练的权重，可以自己另外加载，但是这里就直接用官方提供的了。在第一次运行的时候会自动下载相应的模型（例如这里就是vgg16），如果弹出了类似“time out”之类的错误的话请运行多一次试试看。通常运行多几次就可以成功将模型下载下来。

然后需要确定的就是模型的结构，只需要：

feature = torch.nn.Sequential(*list(model.children())[:])
print(feature)

例如vgg16的输出是：

Sequential(
  (0): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (1): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace)
    (2): Dropout(p=0.5)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace)
    (5): Dropout(p=0.5)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

基本上可以看出总体分为两个部分，这两个部分对应的名字可以用print(model._modules.keys())查询到：

# 下面是输出内容
odict_keys(['features', 'classifier'])

之后可以直接用model.features直接只调用features部分（注意，不是所有torchvision提供的模型都有model.features这个key，不过目前确认vgg和resnet有），而将分类部分抛弃掉。例如torchvision官方例子：

print(input_224.size())            # (1,3,224,224)
output = model.features(input_224) 
print(output.size())               # (1,2048,1,1)

# print(input_448.size())          # (1,3,448,448)
output = model.features(input_448)
# print(output.size())             # (1,2048,7,7)

总体代码

注意，以下内容中的图片为三通道的图片，如果输入的不是三通道的，而是黑白图片，需要先预处理将黑白图片转换成三通道的图片（其实就是将原单通道图片的内容复制三次弄成三通道，使得三个通道的值都一样）。

import os
import numpy as np
 
import torch
import torch.nn
import torchvision.models as models
from torch.autograd import Variable 
import torch.cuda
import torchvision.transforms as transforms
 
from PIL import Image

TARGET_IMG_SIZE = 224
img_to_tensor = transforms.ToTensor()
 
def make_model():
    model=models.vgg16(pretrained=True).features[:28]	# 其实就是定位到第28层，对照着上面的key看就可以理解
    model=model.eval()	# 一定要有这行，不然运算速度会变慢（要求梯度）而且会影响结果
    model.cuda()	# 将模型从CPU发送到GPU,如果没有GPU则删除该行
    return model
    
#特征提取
def extract_feature(model,imgpath):
    model.eval()		# 必须要有，不然会影响特征提取结果
    
    img=Image.open(imgpath)		# 读取图片
    img=img.resize((TARGET_IMG_SIZE, TARGET_IMG_SIZE))
    tensor=img_to_tensor(img)	# 将图片转化成tensor
    tensor=tensor.cuda()	# 如果只是在cpu上跑的话要将这行去掉
    
    result=model(Variable(tensor))
    result_npy=result.data.cpu().numpy()	# 保存的时候一定要记得转成cpu形式的，不然可能会出错
    
    return result_npy[0]	# 返回的矩阵shape是[1, 512, 14, 14]，这么做是为了让shape变回[512, 14,14]
    
if __name__=="__main__":
    model=make_model()
    imgpath='/path/to/img.jpg'
    tmp = extract_feature(model, imgpath)
    print(tmp.shape)	# 打印出得到的tensor的shape
    print(tmp)		# 打印出tensor的内容，其实可以换成保存tensor的语句，这里的话就留给读者自由发挥了

参考

使用pytorch预训练模型分类与特征提取：这篇博文最重要，本博文中代码基本就是在这篇博文的基础上改的
如何从训练好的 PyTorch 模型中提取一幅图像的特征？

---

1. 详情见PyTorch 有哪些坑/bug？中御宅暴君的回答 ↩︎