代码如下：
from PIL import Image
from torchvision import transforms
import os
import torch
import torchvision
import torch.nn.functional as Fclass VGGSim(torch.nn.Module):def __init__(self):super(VGGSim, self).__init__()blocks = []blocks.append(torchvision.models.vgg16(pretrained=True).features[:4].eval())blocks.append(torchvision.models.vgg16(pretrained=True).features[4:9].eval())blocks.append(torchvision.models.vgg16(pretrained=True).features[9:16].eval())blocks.append(torchvision.models.vgg16(pretrained=True).features[16:23].eval())for bl in blocks:for p in bl:p.requires_grad = Falseself.blocks = torch.nn.ModuleList(blocks)self.transform = torch.nn.functional.interpolateself.mean = torch.nn.Parameter(torch.tensor([0.485, 0.456, 0.406]).view(1,3,1,1))self.std = torch.nn.Parameter(torch.tensor([0.229, 0.224, 0.225]).view(1,3,1,1))def forward(self, input, target):if input.shape[1] != 3:input = input.repeat(1, 3, 1, 1)target = target.repeat(1, 3, 1, 1)input = (input-self.mean) / self.stdtarget = (target-self.mean) / self.stdx = inputy = targetres = []for block in self.blocks:x = block(x)y = block(y)x_flat = torch.flatten(x, start_dim=1)y_flat = torch.flatten(y, start_dim=1)similarity = torch.nn.functional.cosine_similarity(x_flat, y_flat)res.append(similarity.cpu().item())# 仅利用VGG最后一层的全局(分类)特征计算余弦相似度# return res[-1]# 或者，利用VGG各Block的特征计算余弦相似度return sum(res)def load_image(path):image = Image.open(path).convert('RGB')image = transforms.Resize([224,224])(image)image = transforms.ToTensor()(image)image = image.unsqueeze(0)return image.cuda()query_image_path = "query.jpeg" # 想要查找的图像
query_image = load_image(query_image_path)
target_image_dir = "cat_images/" # 待搜索的相册
target_images = [os.path.join(target_image_dir, name) for name in os.listdir(target_image_dir)]
vgg_sim = VGGSim().cuda()
scores = []
for path in target_images:target_image = load_image(path)score = vgg_sim(query_image, target_image)scores.append([path, score])
scores.sort(key=lambda x: -x[1])
for i in range(5):print("Top", (i + 1), "similiar =>", scores[i][0].split("/")[-1])
上述代码的核心思想类似于感知损失(Perceptual Loss)，利用VGG提取图像的多级特征，从而比较两张图像之间的相似性。区别在于Perceptual Loss中一般使用MAE，MSE比较特征的距离，而这里的代码使用余弦相似度。
一个例子如下，给定一张狸花的图像(query)如下：
我们希望找到相册中其他狸花的图像：
上述数据集中，编号01到10的为奶牛猫，编号11到20的则为狸花猫。运行代码，结果如下：
Top 1 similiar => 04.jpeg
Top 2 similiar => 20.jpeg
Top 3 similiar => 14.jpeg
Top 4 similiar => 12.jpeg
Top 5 similiar => 15.jpeg
可以看到，检索基本是正确的，20，14，12，15均为狸花猫。04得到最高相似度的原因是其与query的姿势十分相似，且环境也差不多(地板)，这也是另一种层面上的两图像相似。
)