版权声明：本文为博主原创文章，如未特别声明，均默认使用CC BY-SA 3.0许可。 https://blog.csdn.net/Geek_of_CSDN/article/details/84343971
前言
这里的提取图片特征特指从VGG网络的最后一个conv层进行提取。虽然下面代码里面给出的是VGG16作为例子，其实也可以用其他的已经经过训练了的神经网络，包括自己训练的。
代码
模型结构相关基本知识
首先说下加载模型，这里用的是torch官方提供的已经训练好的模型，只需要从torchvision模块导入：
import torchvision.models as models
model = models.vgg16(pretrained=True)
上面的pretrained=True是指使用预训练的权重，可以自己另外加载，但是这里就直接用官方提供的了。在第一次运行的时候会自动下载相应的模型（例如这里就是vgg16），如果弹出了类似“time out”之类的错误的话请运行多一次试试看。通常运行多几次就可以成功将模型下载下来。
然后需要确定的就是模型的结构，只需要：
feature = 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部分，直接将分类部分抛弃掉。
总体代码
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 = 448
img_to_tensor = transforms.ToTensor()
def make_model():
model=models.vgg16(pretrained=True).features[:29] # 其实就是定位到第29层，对照着上面的key看就可以理解
model.cuda() # 将模型从CPU发送到GPU,如果没有GPU则删除该行
return model
#特征提取
def extract_feature(model,imgpath):
model.eval() # 必须要有，不然会影响特征提取结果，因为这句话会影响到模型内dropout层的行为
img=Image.open(imgpath) # 读取图片
img=img.resize((TARGET_IMG_SIZE, TARGET_IMG_SIZE))
tensor=img_to_tensor(img) # 将图片转化成tensor
tensor=tensor.resize_(1,3,224,224) # 因为pytorch需要输入的tensor的shape就是这样，增加一个“1”对提取结果不会产生影响
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 模型中提取一幅图像的特征？