基于clip-gmp-vit-l-14预训练模型的图像检索系统实战：从部署到streamlit可视化

1次阅读

没有评论

共计 2152 个字符，预计需要花费 6 分钟才能阅读完成。

在构建图像检索系统时，使用大规模预训练模型（如 clip-gmp-vit-l-14）往往会遇到以下挑战：

计算资源消耗大：这类模型通常需要 GPU 支持，显存占用高
推理延迟高：单次处理耗时可能达到数百毫秒，影响用户体验
特征管理复杂：需要高效存储和检索百万级特征向量
可视化需求强：业务人员需要直观理解检索结果

原生 PyTorch：
优点：开发简单，支持完整模型功能
缺点：推理效率低，内存占用高
ONNX Runtime：
优点：跨平台，支持硬件加速
缺点：转换过程可能遇到算子不支持问题
TorchScript：
优点：保持 PyTorch 特性，部署简单
缺点：优化程度有限

最终选择 原生 PyTorch+ONNX Runtime混合方案，平衡开发效率与运行性能。

关键优化点：

延迟加载：只在需要时加载模型
批处理：合并多个请求提高吞吐量
混合精度：使用 FP16 减少显存占用
量化：对非关键层使用 INT8 量化

import torch
from transformers import CLIPModel, CLIPProcessor

# 延迟加载单例模式
class ClipWrapper:
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance.model = CLIPModel.from_pretrained("clip-gmp-vit-l-14")
            cls._instance.processor = CLIPProcessor.from_pretrained("clip-gmp-vit-l-14")
            cls._instance.model = cls._instance.model.half().cuda()  # FP16 优化
        return cls._instance

采用 FAISS 进行向量检索，优势：

支持 GPU 加速
内置多种索引算法
支持增量更新

import faiss
import numpy as np

# 创建索引
dim = 512  # clip-gmp-vit-l-14 特征维度
index = faiss.IndexFlatIP(dim)
if torch.cuda.is_available():
    res = faiss.StandardGpuResources()
    index = faiss.index_cpu_to_gpu(res, 0, index)

# 添加特征
features = np.random.rand(1000, dim).astype('float32')
index.add(features)

# 检索
query = np.random.rand(1, dim).astype('float32')
D, I = index.search(query, k=5)  # 返回距离和索引

主要组件：

文件上传区域
结果显示区域
参数调节侧边栏

import streamlit as st
from PIL import Image

st.title('图像检索系统')

# 文件上传
uploaded_file = st.file_uploader("上传查询图片", type=["jpg", "png"])

# 参数设置
threshold = st.sidebar.slider("相似度阈值", 0.0, 1.0, 0.7)
top_k = st.sidebar.number_input("返回结果数", 1, 20, 5)

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption='查询图片', use_column_width=True)

    # 执行检索
    with st.spinner('正在检索...'):
        features = extract_features(image)
        D, I = index.search(features, k=top_k)

    # 显示结果
    for i, (score, idx) in enumerate(zip(D[0], I[0])):
        if score > threshold:
            st.image(result_images[idx], caption=f'结果{i+1} 相似度:{score:.2f}')

测试环境：NVIDIA T4 GPU, 16GB 内存