pytorch深度学习以图搜图

---

前言

最近项目上有一些图像相似性的问题需要研究，之前用传统基于特征点方法还是有一些劣势。想了一下写一篇简易的关于使用神经网络来做图像搜索的文章，图像搜索本质是输入一张图像，从数据库查找到和他最相似的图像并排序返回。最关键的环节就是两张图像相似度的度量。本文方法感觉和孪生神经网络没有什么本质区别，都是输入到同一个网络然后计算相似性，不过本文是将数据库大量图像特征提前计算保存下来，也可以节省计算时间。本文一共三部分代码。

---

一、特征提取网络

可以使用在一些公开数据集(比如imagenet，cifar)上面训练好的模型作为特征提取器，例如vgg16或者resnet这些，都有很好的特征提取能力。如果是特定垂类的数据建议重新训练一下。我这里是重新训练了一个图像二分类模型，backbone使用了华为的G-ghostnet
代码：

# G_ghostnet_extract.py
import numpy as np
import torch
import torch.nn as nn

__all__ = ['regnetx_032']


def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=dilation, groups=groups, bias=False, dilation=dilation)


def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)


class Bottleneck(nn.Module):
    expansion = 1
    __constants__ = ['downsample']

    def __init__(self, inplanes, planes, stride=1, downsample=None, group_width=1,
                 dilation=1, norm_layer=None):
        super(Bottleneck, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        width = planes * self.expansion
        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
        self.conv1 = conv1x1(inplanes, width)
        self.bn1 = norm_layer(width)
        self.conv2 = conv3x3(width, width, stride, width // min(width, group_width), dilation)
        self.bn2 = norm_layer(width)
        self.conv3 = conv1x1(width, planes)
        self.bn3 = norm_layer(planes)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out


class LambdaLayer(nn.Module):
    def __init__(self, lambd):
        super(LambdaLayer, self).__init__()
        self.lambd = lambd

    def forward(self, x):
        return self.lambd(x)


class Stage(nn.Module):

    def __init__(self, block, inplanes, planes, group_width, blocks, stride=1, dilate=False, cheap_ratio=0.5):
        super(Stage, self).__init__()
        norm_layer = nn.BatchNorm2d
        downsample = None
        self.dilation = 1
        previous_dilation = self.dilation
        if dilate:
            self.dilation *= stride
            stride = 1
        if stride != 1 or self.inplanes != planes:
            downsample = nn.Sequential(
                conv1x1(inplanes, planes, stride),
                norm_layer(planes),
            )

        self.base = block(inplanes, planes, stride, downsample, group_width,
                          previous_dilation, norm_layer)
        self.end = block(planes, planes, group_width=group_width,
                         dilation=self.dilation,
                         norm_layer=norm_layer)

        group_width = int(group_width * 0.75)
        raw_planes = int(planes * (1 - cheap_ratio) / group_width) * group_width
        cheap_planes = planes - raw_planes
        self.cheap_planes = cheap_planes
        self.raw_planes = raw_planes

        self.merge = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(planes + raw_planes * (blocks - 2), cheap_planes,
                      kernel_size=1, stride=1, bias=False),
            nn.BatchNorm2d(cheap_planes),
            nn.ReLU(inplace=True),
            nn.Conv2d(cheap_planes, cheap_planes, kernel_size=1, bias=False),
            nn.BatchNorm2d(cheap_planes),
            #             nn.ReLU(inplace=True),
        )
        self.cheap = nn.Sequential(
            nn.Conv2d(cheap_planes, cheap_planes,
                      kernel_size=1, stride=1, bias=False),
            nn.BatchNorm2d(cheap_planes),
            #             nn.ReLU(inplace=True),
        )
        self.cheap_relu = nn.ReLU(inplace=True)

        layers = []
        downsample = nn.Sequential(
            LambdaLayer(lambda x: x[:, :raw_planes])
        )

        layers = []
        layers.append(block(raw_planes, raw_planes, 1, downsample, group_width,
                            self.dilation, norm_layer))
        inplanes = raw_planes
        for _ in range(2, blocks - 1):
            layers.append(block(inplanes, raw_planes, group_width=group_width,
                                dilation=self.dilation,
                                norm_layer=norm_layer))

        self.layers = nn.Sequential(*layers)

    def forward(self, input):
        x0 = self.base(input)

        m_list = [x0]
        e = x0[:, :self.raw_planes]
        for l in self.layers:
            e = l(e)
            m_list.append(e)
        m = torch.cat(m_list, 1)
        m = self.merge(m)

        c = x0[:, self.raw_planes:]
        c = self.cheap_relu(self.cheap(c) + m)

        x = torch.cat((e, c), 1)
        x = self.end(x)
        return x


class RegNet(nn.Module):

    def __init__(self, block, layers, widths, num_classes=1000, zero_init_residual=True,
                 group_width=1, replace_stride_with_dilation=None,
                 norm_layer=None):
        super(RegNet, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self._norm_layer = norm_layer

        self.inplanes = 32
        self.dilation = 1
        if replace_stride_with_dilation is None:
            # each element in the tuple indicates if we should replace
            # the 2x2 stride with a dilated convolution instead
            replace_stride_with_dilation = [False, False, False, False]
        if len(replace_stride_with_dilation) != 4:
            raise ValueError("replace_stride_with_dilation should be None "
                             "or a 4-element tuple, got {}".format(replace_stride_with_dilation))
        self.group_width = group_width
        self.conv1 = nn.Conv2d(1, self.inplanes, kernel_size=3, stride=2, padding=1,
                               bias=False)
        self.bn1 = norm_layer(self.inplanes)
        self.relu = nn.ReLU(inplace=True)

        self.layer1 = self._make_layer(block, widths[0], layers[0], stride=2,
                                       dilate=replace_stride_with_dilation[0])

        self.inplanes = widths[0]
        if layers[1] > 2:
            self.layer2 = Stage(block, self.inplanes, widths[1], group_width, layers[1], stride=2,
                                dilate=replace_stride_with_dilation[1], cheap_ratio=0.5)
        else:
            self.layer2 = self._make_layer(block, widths[1], layers[1], stride=2,
                                           dilate=replace_stride_with_dilation[1])

        self.inplanes = widths[1]
        self.layer3 = Stage(block, self.inplanes, widths[2], group_width, layers[2], stride=2,
                            dilate=replace_stride_with_dilation[2], cheap_ratio=0.5)

        self.inplanes = widths[2]
        if layers[3] > 2:
            self.layer4 = Stage(block, self.inplanes, widths[3], group_width, layers[3], stride=2,
                                dilate=replace_stride_with_dilation[3], cheap_ratio=0.5)
        else:
            self.layer4 = self._make_layer(block, widths[3], layers[3], stride=2,
                                           dilate=replace_stride_with_dilation[3])
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.dropout = nn.Dropout(0.2)
        self.fc = nn.Linear(widths[-1] * block.expansion, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        # Zero-initialize the last BN in each residual branch,
        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677

    #         if zero_init_residual:
    #             for m in self.modules():
    #                 if isinstance(m, Bottleneck):
    #                     nn.init.constant_(m.bn3.weight, 0)

    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if dilate:
            self.dilation *= stride
            stride = 1
        if stride != 1 or self.inplanes != planes:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes, stride),
                norm_layer(planes),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample, self.group_width,
                            previous_dilation, norm_layer))
        self.inplanes = planes
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes, group_width=self.group_width,
                                dilation=self.dilation,
                                norm_layer=norm_layer))
        return nn.Sequential(*layers)

    def _forward_impl(self, x):
        # See note [TorchScript super()]
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        # x = x.cpu()  # 如果用cuda取消这一行的注释,否则下面numpy计算会报错
        LA = np.linalg.norm(x[0])  # 二范数l2 = sqrt(x1^2+x2^2+x3^)
        norm_feat = np.array(x[0]) / LA  # 01之间
        return norm_feat
        # x = self.dropout(x)
        # x = self.fc(x)
        # return x

    def forward(self, x):
        return self._forward_impl(x)


def regnetx_032(*kwargs):
    return RegNet(Bottleneck, [2, 6, 15, 2], [96, 192, 432, 1008], group_width=48, num_classes=2, *kwargs)

这部分代码基本与开源一致改动就几行，_forward_impl函数flatten后直接获取到特征向量输出，regnetx_032是传参主函数num_classes=2。
训练时这部分代码还是按照正常的流程来做：

x = self.avgpool(x)
x = torch.flatten(x, 1)
x = self.dropout(x)
x = self.fc(x)
return x

二、数据库图像特征提取

计算多个图像的特征并保存成npz文件。
代码：

# save_features.py
import torch
from torchvision import transforms
from PIL import Image
from G_ghostnet_extract import regnetx_032
import numpy as np
from glob import glob


def predict(savefileName, imgs):
    device = torch.device('cpu')
    torch.set_grad_enabled(False)
    model = regnetx_032()
    model.load_state_dict(torch.load('./ghostnet-10-regular.pth'), False)  # 加载模型
    model.eval()
    model = model.to(device)
    transform_data = transforms.Compose([transforms.ToTensor()])
    featuresList = []
    namesList = []
    img_list = sorted(glob(imgs+"/*"))
    for img_path in img_list:
        img = Image.open(img_path).convert("L")  # 我这里是训练的灰度图
        img = img.resize((112, 112), 0)  # 输入大小我这里是112,与训练一致就行
        img = transform_data(img).unsqueeze(0)
        img = img.to(device)
        output = model(img)  # 输出特征
        featuresList.append(output)
        namesList.append(img_path.split('/')[-1])
    featuresList = np.array(featuresList)
    namesList = np.string_(namesList)  # np.string_得到的是bytes类型,比str更省空间
    np.savez(savefileName, featuresList, namesList)


if __name__ == "__main__":
    savefileName = "features"  # 保存特征的名称
    imgs = "imgs"  # 输入图像文件夹,这个文件夹下放了很多图片
    predict(savefileName, imgs)

执行之后会在当前路径得到features.npz文件，这就是imgs文件夹里面所有图像保存的特征。

三、特征比对计算

此时输入一个图像获取特征，再与加载进来的图像特征文件计算相似性返回结果
代码：

# load_features.py
import time
import numpy as np
import torch
from torchvision import transforms
from PIL import Image
from G_ghostnet_extract import regnetx_032


def cal_feature(img_path):
    transform_data = transforms.Compose([transforms.ToTensor()])
    img = Image.open(img_path).convert("L")
    img = img.resize((112, 112), 0)
    img = transform_data(img).unsqueeze(0)
    img = img.to(device)
    output = model(img)  # 输出特征
    return output


if __name__ == "__main__":
    device = torch.device('cpu')
    torch.set_grad_enabled(False)
    model = regnetx_032()
    model.load_state_dict(torch.load('./ghostnet-10-regular.pth'), False)
    model.eval()
    model = model.to(device)
    start = time.time()
    savefileName = "features"
    readNpy = np.load(savefileName+".npz")  # 加载之前计算的图像集特征文件
    # np.savez会自动命名参数arr_0, arr_1,也可以自己传入
    features = readNpy["arr_0"]  # 特征向量
    names = readNpy["arr_1"]  # 图像名称
    imgFeature = cal_feature("./imgs/1.tif")  # 输入待查找的图像获取特征
    scores = np.dot(imgFeature, features.T)  # 计算得分0-1之间,越靠近1越相似
    # 得分从大到小排序
    sortID = np.argsort(scores)[::-1]
    sort_score = scores[sortID]
    returnNum = 200  # 搜索出200张相似度最高的图片
    imgList = []
    for i, index in enumerate(sortID[0:returnNum]):
        imgList.append(names[index].decode('utf-8'))
        print("name:{}, score:{}".format(names[index].decode('utf-8'), sort_score[i]))
    print("花费时间", time.time()-start)
    print("返回top %d图像排序list" % returnNum, imgList)

---

从imgs取一张图1.tif测试，两张图一样所以计算出来的score为1，后面的相似度依次排序

返回top200，230张图像的数据cpu计算时间大概80ms左右，感觉还是不错的。

---

总结

本文主要写了一下如何使用神经网络做图像搜索，其本质是和批量特征比对得到相似性，应该也可以用于人脸识别，商品搜索等内容，因为没做过上述实际场景的项目就不多描述，如果写的有错误的地方望指正交流，感谢！！！