使用 PyTorch 的图像向量

一篇关于从头构建第一个版本的 Medium 文章：https://becominghuman.ai/extract-a-feature-vector-for-any-image-with-pytorch-9717561d1d4c

图像嵌入的应用：

• 推荐系统的排序
• 将图像聚类到不同类别
• 分类任务
• 图像压缩

可用模型

模型名称	返回向量长度
Resnet-18	512
Resnet-34	512
Resnet-50	2048
Resnet-101	2048
Resnet-152	2048
Alexnet	4096
Vgg-11	4096
Vgg-13	4096
Vgg-16	4096
Vgg-19	4096
Densenet121	1024
Densenet161	2208
Densenet169	1664
Densenet201	1920
efficientnet_b0	1280
efficientnet_b1	1280
efficientnet_b2	1408
efficientnet_b3	1536
efficientnet_b4	1792
efficientnet_b5	2048
efficientnet_b6	2304
efficientnet_b7	2560

安装

已在 Python 3.6 和 torchvision 0.11.0（夜间版，2021-09-25）上测试通过。

需要 PyTorch：http://pytorch.org/

conda install -c pytorch-nightly torchvision

pip install img2vec_pytorch

运行测试

python -m img2vec_pytorch.test_img_to_vec

作为库使用 img2vec

from img2vec_pytorch import Img2Vec
from PIL import Image

# 使用 GPU 初始化 Img2Vec
img2vec = Img2Vec(cuda=True)

# 读取一张图像（RGB 格式）
img = Image.open('test.jpg')
# 从 img2vec 获取向量，返回为 torch FloatTensor
vec = img2vec.get_vec(img, tensor=True)
# 或者提交一个列表
vectors = img2vec.get_vec(list_of_PIL_images)

运行示例还需要：

• Pillow：pip install Pillow
• Sklearn pip install scikit-learn

运行示例

git clone https://github.com/christiansafka/img2vec.git

cd img2vec/example

python test_img_similarity.py

预期输出

您想查询哪张图片的相似度？
cat.jpg
0.72832 cat2.jpg
0.641478 catdog.jpg
0.575845 face.jpg
0.516689 face2.jpg

您想查询哪张图片的相似度？
face2.jpg
0.668525 face.jpg
0.516689 cat.jpg
0.50084 cat2.jpg
0.484863 catdog.jpg

试试加入你自己的照片吧！

Img2Vec 参数

cuda = (True, False)   # 是否在 GPU 上运行？     默认：False
model = ('resnet-18', 'efficientnet_b0', 等)   # 使用哪个模型？     默认：'resnet-18'

高级用户

---

只读文件系统

如果您在只读环境中运行此库（例如 Docker 容器），请指定一个可写目录，以便应用程序可以存储预训练模型。

export TORCH_HOME=/tmp/torch

其他参数

layer = 'layer_name' 或 int   # 对于高级用户，指定从模型的哪一层提取输出。   默认：'avgpool'
layer_output_size = int   # 您选择的层的输出大小
gpu = (0, 1, 等)   # 使用哪块 GPU？     默认：0

Resnet-18

默认值：(layer = 'avgpool', layer_output_size = 512)
Layer 参数必须是一个表示以下层名称的字符串：

conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
bn1 = nn.BatchNorm2d(64)
relu = nn.ReLU(inplace=True)
maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
layer1 = self._make_layer(block, 64, layers[0])
layer2 = self._make_layer(block, 128, layers[1], stride=2)
layer3 = self._make_layer(block, 256, layers[2], stride=2)
layer4 = self._make_layer(block, 512, layers[3], stride=2)
avgpool = nn.AvgPool2d(7)
fc = nn.Linear(512 * block.expansion, num_classes)

Alexnet

默认值：(layer = 2, layer_output_size = 4096)
Layer 参数必须是一个表示以下层之一的整数：

alexnet.classifier = nn.Sequential(
            7. nn.Dropout(),                  < - 输出大小为 9216
            6. nn.Linear(256 * 6 * 6, 4096),  < - 输出大小为 4096
            5. nn.ReLU(inplace=True),         < - 输出大小为 4096
            4. nn.Dropout(),		      < - 输出大小为 4096
            3. nn.Linear(4096, 4096),	      < - 输出大小为 4096
            2. nn.ReLU(inplace=True),         < - 输出大小为 4096
            1. nn.Linear(4096, num_classes),  < - 输出大小为 4096
        )

Vgg

默认值：(layer = 2, layer_output_size = 4096)

vgg.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, num_classes),
        )

Densenet

默认值：(layer = 1 from features, layer_output_size = 1024)

densenet.features = nn.Sequential(OrderedDict([
	('conv0', nn.Conv2d(3, num_init_features, kernel_size=7, stride=2,
						padding=3, bias=False)),
	('norm0', nn.BatchNorm2d(num_init_features)),
	('relu0', nn.ReLU(inplace=True)),
	('pool0', nn.MaxPool2d(kernel_size=3, stride=2, padding=1)),
]))

EfficientNet

默认值：(layer = 1 from features，layer_output_size = 1280 for efficientnet_b0 model)

待办事项

• 基准测试速度和精度
• 添加对输入数据进行微调的功能
• 将文档导出到正常位置

img2vec 快速上手指南

img2vec 是一个基于 PyTorch 的轻量级库，能够将任意图像转换为特征向量（Embedding）。生成的向量可用于图像相似度搜索、推荐系统排序、聚类分析及图像分类等任务。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux, macOS 或 Windows
• Python 版本：推荐 Python 3.6+
• 核心依赖：
  • PyTorch (需预先安装)
  • torchvision (推荐版本 0.11.0+)
  • Pillow (用于图像处理)
  • scikit-learn (运行示例代码时需要)

国内加速建议： 安装 PyTorch 及相关依赖时，推荐使用清华源或阿里源以提升下载速度。

pip install torch torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

您可以通过 pip 直接安装预编译包：

pip install img2vec_pytorch

如果需要运行官方示例代码，请额外安装以下依赖：

pip install Pillow scikit-learn

若需体验最新特性或进行开发测试，可安装 nightly 版本的 torchvision：

conda install -c pytorch-nightly torchvision

基本使用

以下是使用 img2vec 提取图像特征向量的最简代码示例。

1. 初始化与单张图片处理

from img2vec_pytorch import Img2Vec
from PIL import Image

# 初始化 Img2Vec
# cuda=True 表示使用 GPU 加速，若无 GPU 请设为 False 或省略该参数
img2vec = Img2Vec(cuda=True)

# 读取图片 (RGB 格式)
img = Image.open('test.jpg')

# 获取特征向量
# tensor=True 返回 torch FloatTensor 类型，默认为 numpy 数组
vec = img2vec.get_vec(img, tensor=True)

print(f"向量形状：{vec.shape}")

2. 批量处理图片

您也可以直接传入一个 PIL Image 对象列表进行批量处理：

# 假设 list_of_PIL_images 是一个包含多个 PIL Image 对象的列表
vectors = img2vec.get_vec(list_of_PIL_images)

3. 自定义模型与参数

初始化时可指定不同的预训练模型（如 resnet-50, efficientnet_b0 等）：

# 使用 ResNet-50 模型，输出向量长度为 2048
img2vec = Img2Vec(model='resnet-50', cuda=True)

# 或者指定 EfficientNet
# img2vec = Img2Vec(model='efficientnet_b0', cuda=True)

常用模型及其对应的向量长度参考：

• resnet-18: 512
• resnet-50: 2048
• vgg-16: 4096
• efficientnet_b0: 1280

4. 运行官方示例（图像相似度）

如果您克隆了源代码仓库，可以运行内置的相似度测试脚本：

git clone https://github.com/christiansafka/img2vec.git
cd img2vec/example
python test_img_similarity.py

运行后根据提示输入文件名，系统将输出与该图片最相似的其他图片及其相似度得分。