该项目是一个基于 OpenAI 的 CLIP
模型的应用,旨在通过输入文本描述来检索与之相匹配的图片。它结合了 CLIP
模型的强大图像和文本理解能力,使得用户可以通过简单的文本查询来发现感兴趣的图片。为了提高检索效率,项目采用了 Milvus
和 Redis
数据库,并利用 ONNX
进行推理加速。
-
图文检索:用户可以输入文本描述,系统将利用 CLIP 模型从数据库中检索出与之最匹配的图片,并返回结果。
-
索引技术:项目使用 Milvus 数据库来存储图像特征向量,利用其高效的相似度搜索功能进行图像检索。
-
缓存优化:通过 Redis 数据库,将一部分检索结果缓存在内存中,以加速后续的相似度搜索。(可选功能)
-
ONNX 推理加速:通过将 CLIP 模型转换为 ONNX 格式,利用 ONNX Runtime 进行推理加速,提高检索的实时性。
R@1 | R@3 | R@10 | MRR | MAP | NDCG@50 | Precision@50 | |
---|---|---|---|---|---|---|---|
Mini-ImageNet | 98.0 | 98.0 | 98.0 | 0.980 | 0.980 | 0.980 | 0.980 |
ImageNet-val | 93.1 | 94.9 | 95.3 | 0.940 | 0.936 | 0.931 | 0.929 |
R@1 | R@3 | R@10 | MRR | MAP | NDCG@1 | Precision@1 | |
---|---|---|---|---|---|---|---|
MSCOCO-5k | 78.7 | 89.7 | 86.1 | 0.850 | 0.850 | 0.787 | 0.787 |
MSCOCO | 60.7 | 74.7 | 85.6 | 0.694 | 0.693 | 0.607 | 0.607 |
Flickr30K | 93.7 | 98.1 | 99.3 | 0.960 | 0.960 | 0.937 | 0.937 |
在这里提供项目的安装指南:
- 下载安装单机版
milvus
向量数据库 https://milvus.io/docs/install_standalone-docker.md
$ wget https://github.com/milvus-io/milvus/releases/download/v2.2.12/milvus-standalone-docker-compose.yml -O docker-compose.yml
$ sudo docker-compose up -d
$ docker port milvus-standalone 19530/tcp
- 下载安装
Redis
(可选功能)
$ sudo apt-get install redis-server
- 安装所需依赖项:
$ git clone https://github.com/sjy0727/CLIP-Text-Image-Retrieval.git
$ cd ./CLIP-Text-Image-Retrieval
$ pip install -r requirements.txt
-
下载
mini-imagenet
数据集。 -
把
mini-imagenet,MSCOCO,flickr30k
等文件夹与*.py
文件放在同级目录下。
mini-imagenet
├── classes_name.json
├── imagenet_class_index.json
├── images
│ ├── xxx.jpg
│ ├── xxx.jpg
│ └── ...
├── new_train.csv
├── new_val.csv
├── train.json
└── val.json
下载coco2014
数据集: https://cocodataset.org/#download
MSCOCO
├── annotations
│ ├── captions_train2014.json
│ └── captions_val2014.json
├── train2014
│ ├── xxx.jpg
│ └── ...
└── val2014
├── xxx.jpg
└── ...
下载flickr30k
数据集: http://shannon.cs.illinois.edu/DenotationGraph/data/index.html
Flickr30k
├── flickr30k-images
│ ├── xxx.jpg
│ └── ...
└── results_20130124.token
- 本项目将微调,导出onnx,推理相关的参数写在
config.py
中。可通过指定参数来灵活替换模型及数据库等参数。
- 通过
Huggingface
的accelerate
实现单机多卡训练。
$ accelerate launch --multi_gpu \
--mixed_precision=fp16 \
--num_processes=2 \
finetuning.py \
--ep=5 \
--lr=5e-5 \
--bs=64
- 通过运行
embed2db.py
将数据集存入milvus中。
$ python embed2db.py
- 将模型经
PyTorch
权重转为onnx
权重,并用onnxruntime-gpu
实现推理加速。 - 实现了
动态batch
高效利用硬件算力。
$ python export_onnx.py
- 启动
Grdio
界面。
$ python demo.py
- 通过
http://localhost:7860
访问。