img2vec
PyTorch预训练模型图片特征提取库，三行代码将图片转为向量
加载项目详情…
本应用为开源项目，仅供学习研究，请遵守其开源协议。
微信扫码咨询客服
POMOAI — 让开源 AI 即刻运行
PyTorch预训练模型图片特征提取库，三行代码将图片转为向量
加载项目详情…
本应用为开源项目，仅供学习研究，请遵守其开源协议。
想象一下这样的场景：你手里有上万张商品图片，需要找出外观相似的款式；或者你的相册里有几千张照片，想按内容自动归类。传统的图像比对方式——逐像素比较——在这种规模下根本不现实。而 img2vec_pytorch 正是为解决这个痛点而生：它把任意图片"翻译"成一个数字向量，让图片之间的相似度比较变成简单的数学运算。
这个工具的作者是 Christian Safka，他在 2017 年发表于 Medium 的博文中详细描述了从零构建第一个版本的全过程。这篇博文至今仍是理解图像特征提取原理的经典入门材料。img2vec_pytorch 的核心思想来自计算机视觉领域的一个成熟洞察：训练好的图像分类模型（如 ResNet）在处理日常图片时，其隐藏层的激活值天然携带了强大的语义信息——这比直接比较像素值要可靠得多。
在 2017 年那会儿，PyTorch 的 torchvision 还不像今天这样成熟，作者需要手动裁剪模型结构、注册钩子来提取中间层输出。如今 img2vec_pytorch 已经升级为对 torchvision 最新 API 的封装，开发体验大幅提升。
img2vec_pytorch 本质上是一个预训练模型特征提取器。用户传入一张 PIL 图片，工具会：
标准化预处理：将图片缩放到 224x224 像素，按照 ImageNet 数据集的均值（0.485, 0.456, 0.406）和标准差进行归一化。这确保了预训练模型的输入格式正确。
模型加载：从 torchvision.models 加载用户指定的预训练模型（默认为 ResNet-18）。所有模型权重均来自 ImageNet 预训练，不需要额外训练。
Hook 劫持中间层：这是整个工具最精妙的部分。PyTorch 提供了 register_forward_hook() API，允许在模型前向传播过程中"劫持"任意层的输出。img2vec_pytorch 在目标层（如 ResNet 的 avgpool 层，输出 512 维特征向量）注册一个 hook 函数，将该层的输出复制到预先分配的张量中，而不改变模型原本的计算图。Hook 机制是非侵入式的——不需要修改原始模型结构，不需要重新定义 forward 方法。
输出适配：不同模型的输出形状不同。ResNet 系列输出 (batch, 512, 1, 1)，VGG/AlexNet 输出 (batch, 4096)，DenseNet 和 EfficientNet 则输出 (batch, C, 7, 7) 的 4D 张量，需要通过 squeeze 或 mean 操作降维到 2D 特征向量。最终返回 NumPy 数组（默认）或 PyTorch FloatTensor。
from img2vec_pytorch import Img2Vec
from PIL import Image
img2vec = Img2Vec(cuda=True) # GPU 加速
img = Image.open('photo.jpg').convert('RGB')
vec = img2vec.get_vec(img) # 返回 512 维向量（ResNet-18）
支持批量处理，传入 PIL Image 列表即可一次性提取多个向量，这在实际应用中非常实用。
工具支持 5 大系列共 20 种模型，覆盖从轻量到重型的各种需求：
ResNet 系列（残差网络）：resnet-18/34/50/101/152，输出维度从 512 到 2048。ResNet-18 是速度和精度的均衡之选，也是工具的默认模型。
VGG 系列：vgg11/13/16/19，输出 4096 维向量。VGG 结构简单，但参数量大，适合对特征维度有特殊需求的场景。
DenseNet 系列：densenet121/161/169/201，输出 1024 到 2208 维。DenseNet 以密集跳跃连接著称，特征复用效率高。
EfficientNet 系列：efficientnet_b0 到 b7，输出 1280 到 2560 维。EfficientNet 通过复合缩放策略在精度和效率之间取得了出色的平衡，B0-B7 覆盖从边缘设备到服务器端的各种硬件条件。
AlexNet：4096 维，作为经典baseline保留。
用户可以根据速度/精度 tradeoff 灵活选择。对于推荐系统等需要快速检索的场景，ResNet-18 足够；对于高精度图像分类任务，EfficientNet-B5/B7 效果更好。
图1：示例图片 catdog.jpg，可通过 img2vec 提取 512 维特征向量
项目自带了一个完整示例——test_img_similarity.py，演示如何计算图片之间的余弦相似度。以 ResNet-18 提取的 512 维向量为基础，通过 NumPy 的点积运算即可得到相似度分数。测试集中的 cat.jpg 和 cat2.jpg 相似度高达 0.73（满分 1.0），而 cat 和 face 的相似度仅 0.52，语义层面的差异被向量空间准确捕获。
在实际业务中，开发者通常会将提取的向量存入向量数据库（如 Milvus、FAISS、Qdrant），然后用 ANN（近似最近邻）算法实现毫秒级的相似图片检索。这套流程广泛应用于：
整个项目只有两个核心 Python 文件（img_to_vec.py + init.py），加上一个测试文件，代码总量不超过 500 行。这种极简设计让项目的维护成本极低——没有复杂的依赖树，没有冗余的配置文件，所有功能都围绕一个 Img2Vec 类展开。
测试覆盖了所有 20 种模型，每个模型的输出维度都有对应的断言用例，确保版本更新不会引入回归。测试图片放在 example/test_images/ 目录下，包括 cat、cat2、catdog、face、face2 五张标准测试图片。
安装仅依赖 torch、torchvision、numpy 三件套，不引入多余的机器学习框架。对于 Docker 容器环境，通过设置 TORCH_HOME=/tmp/torch 环境变量可以将预训练模型缓存目录指定为可写位置，解决了只读文件系统下的模型缓存问题。
img2vec_pytorch 的定位是轻量级特征提取工具，而非端到端的图像解决方案。它有几个值得注意的局限：
首先，不支持自定义模型微调。所有特征都来自 ImageNet 预训练权重，对于特定领域的图片（如医学影像、工业零件）可能不是最优选择。这类场景下，用户需要在特定数据集上对模型进行迁移学习，而不仅仅是提取特征。
其次，无 Web UI 和 API 服务。这是一个纯库，推理需要用户在 Python 代码中调用。对于需要对外提供 HTTP 接口的场景，用户需要自行用 FastAPI/Flask 包装。
第三，特征维度固定。不同模型输出的向量维度不同（512 到 4096），在构建向量索引时需要考虑维度一致性。
图像特征提取是 CLIP、SAM 等多模态大模型出现之前，计算机视觉领域最基础也最重要的能力之一。img2vec_pytorch 用不到 500 行代码，将这个能力封装成了一个任何 Python 开发者都能轻松使用的工具。它代表了开源社区在"降低 AI 门槛"方向上的持续努力——从需要理解反向传播和 CUDA 编程，到只需要 pip install 加三行代码。
随着 Vision Transformer（ViT）和多模态大模型的崛起，基于 CNN 特征向量的方法在某些场景下面临挑战。但对于资源受限环境、轻量级检索任务、以及对可解释性有要求的场景，传统 CNN 特征依然是值得信赖的基准线。img2vec_pytorch 作为这类方法的代表工具，在向量检索生态中仍占有独特的一席之地。