图像 Embedding（ViT 系列模型）
更新时间：2026-05-18
图像 Embedding（ViT 系列模型）
简介
ImageViTEmbedding基于 Daft + Ray 分布式框架的图像嵌入向量提取组件，利用 ViT 类模型（如 DINOv2）将图像编码为固定维度的向量表示，适用于图像检索、分类、聚类等下游任务。
功能
- 支持通过图像 URL 批量提取 ViT 嵌入向量
- 基于 Ray Actor 实现 GPU 加速推理
- 支持多种 ViT 模型（如
facebook/dinov2-large） - 支持 FP16 精度推理，降低显存占用
参数
| 参数名称 | 类型 | 默认值 | 描述 |
|---|---|---|---|
| image_src_type | str | image_url | 输入图像的格式类型，支持：bos/http 地址(image_url)base64 编码(image_base64)二进制流(image_binary)可选值：["image_url", "image_base64", "image_binary"]默认值："image_url" |
| dtype | str | float32 | 模型推理精度选择：bfloat16: 平衡精度与速度（TPU上更快）float16: 更快的推理速度float32: 最高精度但显存消耗最大可选值：["bfloat16", "float16", "float32"]默认值："float16" |
| batch_size | int | 32 | 批处理大小默认值: 32 |
| model_path | str | /opt/aihc/models | 模型文件存储路径默认值: "/opt/aihc/models" |
| model_name | str | facebook/dinov2-large | 使用的图像向量模型名称可选值: ["google/vit-base-patch16-224-in21k","google/vit-large-patch16-224-in21k","facebook/dinov2-base","facebook/dinov2-large"]默认值: "facebook/dinov2-large" |
| use_cls_token_embedding | bool | true | 是否使用CLS Token特征默认值: True |
| rank | int | 0 | 指定使用的GPU设备编号（多卡环境有效）。例如：0表示第一张GPU，1表示第二张GPU默认值：0 |
输入
| 输入列名 | 说明 | 说明 |
|---|---|---|
| images | 包含图像数据的数组，元素类型可为图像URL、Base64编码或二进制数据 | 图像 URL 地址 |
以 Daft DataFrame 形式传入，列名为 image。
输出
包含特征向量的数组，每个元素为float类型的嵌套数组， 数组维度由模型输出决定
使用示例
Plain Text
1from __future__ import annotations
2
3import os
4import daft
5from daft import col
6from daft.aihc.common.udf import aihc_udf
7from daft.aihc.functions.image.embedding.image_vit_embedding import ImageViTEmbedding
8
9if __name__ == "__main__":
10 if os.getenv("DAFT_RUNNER", "native") == "ray":
11 import ray
12 ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
13 daft.set_runner_ray()
14 daft.set_execution_config(actor_udf_ready_timeout=6000, min_cpu_per_task=0)
15
16 image_src_type = "image_url"
17 batch_size = 64
18 model_path = os.getenv("MODEL_PATH", "/opt/aihc/models")
19 model_name = "facebook/dinov2-large"
20 dtype = "float16"
21 use_cls_token_embedding = True
22 rank = 0
23 num_gpus = 1
24
25 samples = {
26 "image": [
27 "https://{bucket}.bj.bcebos.com/image.png",
28 ]
29 }
30
31 ds = daft.from_pydict(samples)
32 ds = ds.with_column(
33 "embedding",
34 aihc_udf(
35 ImageViTEmbedding,
36 construct_args={
37 "image_src_type": image_src_type,
38 "batch_size": batch_size,
39 "model_path": model_path,
40 "model_name": model_name,
41 "dtype": dtype,
42 "use_cls_token_embedding": use_cls_token_embedding,
43 "rank": rank,
44 },
45 num_gpus=num_gpus,
46 batch_size=1,
47 )(col("image")),
48 )
49
50 ds.show()
51#╭────────────────────────────────┬────────────────────────────────╮
52#│ image ┆ embedding │
53#│ --- ┆ --- │
54#│ String ┆ List[Float32] │
55#╞════════════════════════════════╪════════════════════════════════╡
56#│ https://{bucket}.bj.bcebos.com/┆ [0.00059747696, -0.02935791, … │
57#╰────────────────────────────────┴────────────────────────────────╯评价此篇文章