使用pytorch预训练模型VGG19提取图像特征, 得到图像embedding
pytorch中的VGG19预训练模型, 最后一层输出是1000维的图像分类结果, 但是如果我们只想要模型中某一层的输出特征, 比如全连接层的4096维度的特征, 要如何提取呢? 本文解决这个问题.
本文内容参考:
https://zhuanlan.zhihu.com/p/105703821
https://blog.csdn.net/Geek_of_/article/details/84343971
https://discuss.pytorch.org/t/how-to-extract-features-of-an-image-from-a-trained-model/119/2
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 # 如果联网使用pytorch的预训练模型, # 将pretrained设置为True, 就会自动下载vgg19的模型放在本地缓存中. vgg_model = models.vgg19(pretrained=True) # 如果使用下载到本地的预训练模型, pretrained默认为False, # 则需要提供本地的模型路径, 并使用load_state_dict加载. # vgg_model = models.vgg19() # pre = torch.load('/XXXX/vgg19-dcbb9e9d.pth') # vgg_model.load_state_dict(pre)代码如下（示例）：
# 查看模型整体结构 structure = torch.nn.Sequential(*list(vgg_model.children())[:]) print(structure) # 查看模型各部分名称 print('模型各部分名称', vgg_model._modules.keys())输出结构如下:
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整体结构分为三大部分'features', 'avgpool', 和 'classifier', 为了得到4096维度的特征, 我们可以去掉classifier层的最后两层(5),(6). 下面介绍如何实现去掉这两层.
上述获取模型结构的代码可以添加参数得到模型的每个部分:
# 获取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部分的基础上, 再切割模型, 得到输出维度为4096的子模型. new_classifier = torch.nn.Sequential(*list(vgg_model.children())[-1][:5]) 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) )
从上述输出可以看出, 与vgg19模型的原始classifier相比, new_classifier的模型结构已经去掉了最后一层全连接层和一个Dropout层, 能够得到输出维度为4096的特征了.
所以下面, 我们就用new_classifier替换vgg19原始模型中的分类器( classifier )部分, 保留vgg19模型前面的features和avgpool不变.代码如下:
# 得到vgg19原始模型, 输出图像维度是1000的. vgg_model_1000 = models.vgg19(pretrained=True) # 下面三行代码功能是:得到修改后的vgg19模型. # 具体实现是: 去掉vgg19模型的最后一个全连接层, 使输出图像维度是4096. vgg_model_4096 = models.vgg19(pretrained=True) # 获取原始模型中去掉最后一个全连接层的classifier. new_classifier = torch.nn.Sequential(*list(vgg_model_4096.children())[-1][:5]) # 替换原始模型中的classifier. vgg_model_4096.classifier = new_classifier # 获取和处理图像 data_dir = '/mnt/image_test.jpg' im = Image.open(data_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得到图像特征表示. image_feature_1000 = vgg_model_1000(im).data[0] image_feature_4096 = vgg_model_4096(im).data[0] print('dim of vgg_model_1000: ', image_feature_1000.shape) 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，目前不知道原因．
