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

学习笔记 同时被 2 个专栏收录
47 篇文章 1 订阅
5 篇文章 0 订阅

注意

目前(2019年一月)因为torchvision提供的VGG网络没有训练完全1,不建议使用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?中御宅暴君的回答 ↩︎

©️2021 CSDN 皮肤主题: 技术黑板 设计师:CSDN官方博客 返回首页
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、C币套餐、付费专栏及课程。

余额充值