使用pytorch中预训练模型VGG19获取图像特征，得到图像embedding

前言

使用pytorch预训练模型VGG19提取图像特征, 得到图像embedding

pytorch中的VGG19预训练模型, 最后一层输出是1000维的图像分类结果, 但是如果我们只想要模型中某一层的输出特征, 比如全连接层的4096维度的特征, 要如何提取呢? 本文解决这个问题.

---

本文内容参考:

pytorch预训练模型的下载地址以及解决下载速度慢的方法 - 知乎

Pytorch：利用预训练好的VGG16网络提取图片特征_Geek_of_CSDN的博客-CSDN博客_pytorch vgg16提取特征

How to extract features of an image from a trained model - #2 by apaszke - PyTorch Forums

一、导入pytorch中vgg19预训练模型

pytorch中vgg预训练模型下载地址:

model_urls = {
    'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth',
    'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth',
    'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth',
    'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth',
    'vgg11_bn': 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth',
    'vgg13_bn': 'https://download.pytorch.org/models/vgg13_bn-abd245e5.pth',
    'vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth',
    'vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth',
}

import torch
import torchvision.models as models

# 参数pretrained默认为False,意味着得到一个随机初始化参数的vgg19模型。 
# vgg_model = models.vgg19()
# 可以通过提供本地参数文件的路径, 并使用load_state_dict加载，
# 得到一个参数是预训练好的vgg19模型。
# pre_file = torch.load('/XXXX/vgg19-dcbb9e9d.pth')
# vgg_model.load_state_dict(pre_file)

# 如果将pretrained设置为True, 意味着直接得到一个加载了预训练参数的vgg19模型。
# 就会自动下载vgg19的参数文件并放在本地缓存中。所以不用提供本地参数文件的路径。
vgg_model = models.vgg19(pretrained=True)

二、使用vgg19得到4096维度的图像特征

1.查看VGG19的模型结构

查看VGG19模型结构的代码如下（示例）：

# 查看模型整体结构
structure = torch.nn.Sequential(*list(vgg_model.children())[:])
print(structure)

# 查看模型各部分名称
print('模型各部分名称', vgg_model._modules.keys())

输出的VGG19模型结构如下:

Sequential(
  (0): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (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=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (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=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (17): ReLU(inplace=True)
    (18): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (19): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (24): ReLU(inplace=True)
    (25): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (26): ReLU(inplace=True)
    (27): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (31): ReLU(inplace=True)
    (32): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (33): ReLU(inplace=True)
    (34): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (35): ReLU(inplace=True)
    (36): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (1): AdaptiveAvgPool2d(output_size=(7, 7))
  (2): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)
模型各部分名称: odict_keys(['features', 'avgpool', 'classifier'])

上述输出显示, vgg19整体结构分为三大部分：

1. 'features'：上面输出的VGG19模型结构中的第一个Sequential，包含（0）-（36）层；
2. 'avgpool'：VGG19模型结构的第二个部分AdaptiveAvgPool2d；
3. 'classifier'：VGG19模型结构的最后一个部分Sequential，包含（0）-（6）层。

我们可以看出来，最后classifier的（6）Linear 层将特征维度从4096降低到了1000维，所以为了得到图像的4096维的特征, 我们需要去掉classifier层的最后一层线性层(6)。 （实际上（3）Linear 操作之后的特征维度就是4096了，所以根据个人需求，你可以选择是要（3）还是（4）或者（5）层之后的输出，它们的输出维度都是4096）

下面介绍如何得到去掉最后一层（6）之后的VGG19模型，也就是得到（5）层的模型输出结果。（如果想要（3）层的输出结果，那么就去掉（4）（5）（6）层）。

2.修改模型结构

上面的获取VGG19模型整体结构的代码，通过添加参数可以得到VGG19模型的各个部分.

所以我们可以先拿到VGG19的第三部分classifier，然后去掉它的最后一层（6）Linear，得到一个新的分类器new_classifier，该分类器的输出维度就是4096了。然后再将新的分类器和VGG19原始模型的第一部分features，第二部分avgpool拼接起来就是一个完整的VGG19模型且输出的特征为4096维度。

# 获取vgg19模型的第一个Sequential, 也就是features部分.
features = torch.nn.Sequential(*list(vgg_model.children())[0])
print('features of vgg19: ', features)

# 获取vgg19模型的最后一个Sequential, 也就是classifier部分.
classifier = torch.nn.Sequential(*list(vgg_model.children())[-1])
print('classifier of vgg19: ', classifier)

# 在获取到最后一个classifier部分的基础上, 再切割模型, 去掉最后一层.
new_classifier = torch.nn.Sequential(*list(vgg_model.children())[-1][:6])
print('new_classifier: ', new_classifier)

上述代码输出为: (为了减少文章长度, 我在下面输出结果里面将features模型的(2)-(33)省略了)

features of vgg19:  Sequential(
  (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (1): ReLU(inplace=True)
  (2).......................................................................(33)

  (34): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (35): ReLU(inplace=True)
  (36): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)
classifier of vgg19:  Sequential(
  (0): Linear(in_features=25088, out_features=4096, bias=True)
  (1): ReLU(inplace=True)
  (2): Dropout(p=0.5, inplace=False)
  (3): Linear(in_features=4096, out_features=4096, bias=True)
  (4): ReLU(inplace=True)
  (5): Dropout(p=0.5, inplace=False)
  (6): Linear(in_features=4096, out_features=1000, bias=True)
)
new_classifier :  Sequential(
  (0): Linear(in_features=25088, out_features=4096, bias=True)
  (1): ReLU(inplace=True)
  (2): Dropout(p=0.5, inplace=False)
  (3): Linear(in_features=4096, out_features=4096, bias=True)
  (4): ReLU(inplace=True)
  (5): Dropout(p=0.5, inplace=False)
)

从上述输出可以看出, 与vgg19模型的原始classifier相比, new_classifier的模型结构已经去掉了最后一个全连接层（6）Linear 了, 所以输出维度是4096了. 

3.得到4096维图像特征的完整代码

综上所述，我们用new_classifier替换vgg19原始模型中的分类器( classifier )部分，就得到了输出维度是4096的VGG19模型，代码如下:

import torch
import torchvision.models as models
from torchvision import transforms
from PIL import Image

# 获取vgg19原始模型, 输出图像维度是1000.
vgg_model_1000 = models.vgg19(pretrained=True)

# 下面三行代码功能是:得到修改后的vgg19模型.
# 具体实现是: 去掉vgg19原始模型的第三部分classifier的最后一个全连接层, 
# 用新的分类器替换原始vgg19的分类器，使输出维度是4096.
vgg_model_4096 = models.vgg19(pretrained=True)
new_classifier = torch.nn.Sequential(*list(vgg_model_4096.children())[-1][:6])
vgg_model_4096.classifier = new_classifier

# 获取和处理图像
image_dir = '/mnt/image_test.jpg'
im = Image.open(image_dir)
trans = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
im = trans(im)
im.unsqueeze_(dim=0)

# 使用vgg19得到图像特征.
# 原始vgg19模型
image_feature_1000 = vgg_model_1000(im).data[0]
print('dim of vgg_model_1000: ', image_feature_1000.shape)

# 修改后的vgg19模型
image_feature_4096 = vgg_model_4096(im).data[0]
print('dim of vgg_model_4096: ', image_feature_4096.shape)

输出结果如下:

dim of vgg_model_1000:  torch.Size([1000])
dim of vgg_model_4096:  torch.Size([4096])

可以看出, 没有修改的vgg19模型输出维度是1000, 修改后的vgg19模型输出图像维度是4096.

---

总结

花了一天多, 解决这个问题, 虽然很费时费力, 但是感觉是有意义的.

扩展

如果不需要得到classifier部分的输出结果, 而是只需要原始vgg19模型的第一个部分features部分的输出结果, 有一种更简单的方法, 代码如下:

# 获取vgg19原始模型的features部分的前34个结构, 得到新的vgg_model模型.
vgg_model = models.vgg19(pretrained=True).features[:34]

# 但是下面的代码只能得到classifier部分的前40个, 
# 而不能得到包含features及avgpool及classifier的一共前40个结构.
# 所以这个方法不能实现输出4096维度图像特征的目标．
# vgg_model = models.vgg19(pretrained=True).classifier[:40]

这种方法只使用于联网加载的vgg19模型（即设置了pretrained=True的模型），不适用于使用了本地vgg19模型的vgg_model，目前不知道原因．