DreamSim：利用合成数据学习人类视觉相似性的新维度

项目页面 | 论文 | Bibtex

Stephanie Fu* $^{1}$, Netanel Tamir* $^{2}$, Shobhita Sundaram* $^{1}$, Lucy Chai $^1$, Richard Zhang $^3$, Tali Dekel $^2$, Phillip Isola $^1$.
(*贡献相等，顺序由随机种子决定)
$^1$ MIT，$^2$ 魏茨曼科学研究所，$^3$ Adobe Research。

摘要

当前用于感知图像相似性的度量方法主要基于像素和图像块。这些度量能够比较图像的低级颜色和纹理特征，但无法捕捉布局、姿态、语义内容等中级差异。而使用图像级嵌入（如 DINO 和 CLIP）的模型则能反映高级别和语义层面的判断，却未必与人类对更细粒度属性的感知完全一致。

DreamSim 是一种新的感知图像相似性度量方法，它弥合了“低级”度量（如 LPIPS、PSNR、SSIM）与“高级”度量（如 CLIP）之间的差距。我们的模型通过拼接 CLIP、OpenCLIP 和 DINO 的嵌入向量进行训练，并进一步在人类感知判断数据上进行微调。我们收集了约 2 万组由扩散模型生成的图像三元组的感知判断数据。实验表明，与现有度量相比，DreamSim 能更好地匹配人类的相似性判断，可用于图像检索等下游任务。

🚀 最新更新

2024年10月14日： 我们发布了 DreamSim 的 4 种新变体！这些新检查点包括：

• 以 DINOv2 B/14 和 SynCLR B/16 作为骨干网络；
• 使用原始对比损失同时在 CLS 特征和密集特征上训练的 DINOv2 B/14 和 DINO B/16。
  • 若要从同时利用 CLS 和 patch 特征训练的模型变体中提取特征，请在调用 dreamsim() 时设置 use_patch_model=True。

这些模型（以及原有模型）将在 我们最新的 NeurIPS 2024 论文《感知对齐何时有益于视觉表征？》 中得到进一步评估。

研究发现，我们的感知对齐表征在多种标准计算机视觉任务上均优于基线模型，包括语义分割、深度估计、目标计数、实例检索以及检索增强生成等。这些结果表明，感知对齐可以作为学习通用视觉表征的一个有效目标。更多详情请参阅论文及我们的 博客文章。

以下是更新后的 NIGHTS 性能表：

	NIGHTS - 验证集	NIGHTS - 测试集
ensemble	96.9%	96.2%
dino_vitb16	95.6%	94.8%
open_clip_vitb32	95.6%	95.3%
clip_vitb32	94.9%	93.6%
dinov2_vitb14	94.9%	95.0%
synclr_vitb16	96.0%	95.9%
dino_vitb16 (patch)	94.9%	94.8%
dinov2_vitb14 (patch)	95.5%	95.1%

2024年9月14日： 我们发布了与 peft>=0.2.0 兼容的集成模型和单分支 DreamSim 模型的新版本。

此外，我们还发布了完整的 10 万条（未过滤）NIGHTS 数据集以及 JND（刚刚可察觉差异）投票数据。

目录

• 要求
• 安装
• 使用
  • 快速入门
  • 单分支模型
  • 特征提取
  • 图像检索
  • 感知损失函数
• NIGHTS 数据集
• 实验
• 引用

要求

• Linux
• Python 3

安装

选项 1： 使用 pip 安装：

pip install dreamsim

该包用于导入和使用 DreamSim 模型。

选项 2： 克隆我们的仓库并安装依赖项。 这是运行我们的训练/评估脚本所必需的。

python3 -m venv ds
source ds/bin/activate
pip install -r requirements.txt
export PYTHONPATH="$PYTHONPATH:$(realpath ./dreamsim)"

若使用 conda 安装：

conda create -n ds
conda activate ds
conda install pip # 用 `which pip` 命令确认
pip install -r requirements.txt
export PYTHONPATH="$PYTHONPATH:$(realpath ./dreamsim)"

使用

有关以下用例的详细示例，请参阅我们的 Colab 演示。

快速入门：感知相似性度量

基本用法是测量两张图像之间的感知距离。分数越高表示越不同，分数越低表示越相似。

以下代码片段即为所需全部内容。首次运行 dreamsim 时，它会自动下载模型权重。默认模型设置在 ./dreamsim/config.py 中指定。若要从同时利用 CLS 和 patch 特征训练的模型变体中提取特征，请在调用 dreamsim() 时设置 use_patch_model=True。

from dreamsim import dreamsim
from PIL import Image

device = "cuda"
model, preprocess = dreamsim(pretrained=True, device=device)

img1 = preprocess(Image.open("img1_path")).to(device)
img2 = preprocess(Image.open("img2_path")).to(device)
distance = model(img1, img2) # 模型接受 [0, 1] 范围内的 RGB 图像，尺寸为 batch_sizex3x224x224

要使用示例图像运行，请执行 demo.py。该脚本应输出距离（0.4453，0.2756）。

单分支模型

默认情况下，DreamSim 使用 CLIP、DINO 和 OpenCLIP（均为 ViT-B/16）的集成模型。如果您需要更轻量级的模型，可以使用仅微调单一骨干网络的 单分支 版本。单分支模型的速度比集成模型快约 3 倍。

可用选项按性能排序为：OpenCLIP-ViTB/32、DINO-ViTB/16、CLIP-ViTB/32。要加载单分支模型，请使用 dreamsim_type 参数。例如：

dreamsim_dino_model, preprocess = dreamsim(pretrained=True, dreamsim_type="dino_vitb16")

特征提取

要使用 dreamsim 提取 单张图像的嵌入，请按照以下代码片段所示使用 embed 方法：

img1 = preprocess(Image.open("img1_path")).to("cuda")
embedding = model.embed(img1)

两张图像之间的感知距离是它们嵌入向量之间的余弦距离。如果嵌入向量已经归一化（默认情况下为真），也可以使用 L2 距离。

图像检索

我们的模型可用于图像检索，并可集成到现有的图像检索流程中。下面的代码根据图像与给定查询图像的相似性对一组图像进行排序。

为了加快速度，我们没有直接对每一对图像调用 model(query, image)，而是先使用 model.embed(image) 方法预计算单张图像的嵌入，然后计算这些嵌入之间的余弦距离。

import pandas as pd
from tqdm import tqdm
import torch.nn.functional as F

# 假设 query 是一张示例图像。
# 假设 images 是我们要搜索的一组图像。

# 计算查询图像的嵌入
query_embed = model.embed(preprocess(query).to("cuda"))
dists = {}

# 计算查询图像与每张搜索图像之间的（余弦）距离
for i, im in tqdm(enumerate(images), total=len(images)):
   img_embed = model.embed(preprocess(im).to("cuda"))
   dists[i] = (1 - F.cosine_similarity(query_embed, img_embed, dim=-1)).item()

# 按距离排序并返回结果
df = pd.DataFrame({"ids": list(dists.keys()), "dists": list(dists.values())})
return df.sort_values(by="dists")

感知损失函数

我们的模型可以用作迭代优化中的损失函数（类似于 LPIPS 度量）。以下是关键代码行；完整示例请参阅 Colab。

for i in range(n_iters):
    dist = model(predicted_image, reference_image)
    dist.backward()
    optimizer.step()

NIGHTS（人类测试相似性的新型图像生成）数据集

DreamSim 通过在 NIGHTS 数据集上进行微调来训练。有关数据集结构和创建的详细信息，请参阅 数据集页面。

运行 ./dataset/download_dataset.sh 可以下载并解压 NIGHTS 数据集到 ./dataset/nights 目录下。解压后的数据集大小为 58 GB。

如果遇到大文件下载困难的问题，可以运行 ./dataset/download_chunked_dataset.sh，将 NIGHTS 数据集拆分为 200 个小压缩包下载。该脚本的输出与 download_dataset.sh 完全相同。

（新！）下载完整的 10 万条预过滤的 NIGHTS 数据集

我们仅使用 2 万个一致同意的三元组进行训练和评估，但出于研究目的，我们也发布了全部 10 万个三元组（其中许多投票较少或存在分歧）。运行 ./dataset/download_unfiltered_dataset.sh 可以下载并解压这个未过滤版本的 NIGHTS 数据集到 ./dataset/nights_unfiltered 目录下。解压后的数据集大小为 289 GB。

（新！）下载 JND 数据

运行 ./dataset/download_jnd_dataset.sh 可以下载最小可觉差（JND）投票数据。CSV 文件将被下载到 ./dataset/jnd_votes.csv。请参阅 Colab 中的示例，了解如何加载 JND 实验数据。

使用 Voxel51 演示可视化 NIGHTS 和嵌入：

实验

下载资源

运行 ./training/download_models.sh 可以下载并解压必要的 ViT 检查点（用于 CLIP、OpenCLIP 和 MAE）到 ./models 目录下。

训练

要在数据集上微调感知模型，请运行 ./training/train.py。例如，要使用 LoRA 微调 DINO、CLIP 和 OpenCLIP 的集成模型，可以运行：

python ./training/train.py --config ./configs/train_ensemble_model_lora.yaml

我们在 ./configs 中提供了更多示例配置，包括使用 LoRA 和 MLP 头部进行微调的示例。完整的训练选项列表及说明请参阅 ./training/train.py。

评估

要对数据集上的感知模型进行评估，请运行 ./evaluation/eval_percep.py。例如，要评估一个训练好的检查点与基准模型在数据集上的表现，可以运行：

python ./evaluation/eval_percep.py --config ./configs/eval_ensemble.yaml

完整的评估选项列表及说明请参阅 ./evaluate/eval_percep.py。

引用

如果您认为我们的工作或任何相关材料有用，请引用我们的论文：

@misc{fu2023dreamsim,
      title={DreamSim：利用合成数据学习人类视觉相似性的新维度}, 
      author={Stephanie Fu、Netanel Tamir、Shobhita Sundaram、Lucy Chai、Richard Zhang、Tali Dekel 和 Phillip Isola},
      year={2023},
      eprint={2306.09344},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

@misc{sundaram2024doesperceptualalignmentbenefit,
      title={感知对齐何时有益于视觉表征？}, 
      author={Shobhita Sundaram、Stephanie Fu、Lukas Muttenthaler、Netanel Y. Tamir、Lucy Chai、Simon Kornblith、Trevor Darrell 和 Phillip Isola},
      year={2024},
      eprint={2410.10817},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2410.10817}, 
}

致谢

我们的代码在ViT特征提取方面借鉴了“Deep ViT Features as Dense Visual Descriptors”仓库，并在代码结构上受到UniverSeg仓库的启发。

DreamSim 快速上手指南

DreamSim 是一种新型感知图像相似度度量工具，旨在弥合“低级”指标（如 LPIPS、PSNR）与“高级”语义指标（如 CLIP）之间的差距。它通过结合人类感知判断进行微调，能更准确地捕捉图像在布局、姿态和语义内容上的中层次差异。

环境准备

• 操作系统: Linux
• Python 版本: Python 3
• 硬件建议: 推荐使用 NVIDIA GPU (CUDA) 以获得最佳推理速度，CPU 亦可运行但较慢。
• 前置依赖: Pillow, torch, tqdm 等（安装时会自动处理）。

安装步骤

你可以选择直接通过 pip 安装（推荐用于仅使用模型），或克隆仓库安装（用于运行训练/评估脚本或访问数据集）。

方式一：使用 pip 安装（推荐）

pip install dreamsim

方式二：从源码安装

如果你需要运行示例脚本或访问完整数据集，请克隆仓库并配置环境：

使用 venv:

python3 -m venv ds
source ds/bin/activate
git clone https://github.com/ssundaram21/dreamsim.git
cd dreamsim
pip install -r requirements.txt
export PYTHONPATH="$PYTHONPATH:$(realpath ./dreamsim)"

使用 Conda:

conda create -n ds python=3.8
conda activate ds
git clone https://github.com/ssundaram21/dreamsim.git
cd dreamsim
pip install -r requirements.txt
export PYTHONPATH="$PYTHONPATH:$(realpath ./dreamsim)"

提示: 国内用户若下载依赖较慢，可在 pip install 命令后添加 -i https://pypi.tuna.tsinghua.edu.cn/simple 使用清华镜像源。

基本使用

1. 计算图像感知距离

这是最常用的功能。首次运行时会自动下载预训练权重。分数越高表示差异越大，分数越低表示越相似。

from dreamsim import dreamsim
from PIL import Image

# 设置设备 (cuda 或 cpu)
device = "cuda"

# 加载模型和预处理函数
# 首次运行会自动下载权重
model, preprocess = dreamsim(pretrained=True, device=device)

# 加载并预处理图片
img1 = preprocess(Image.open("img1_path")).to(device)
img2 = preprocess(Image.open("img2_path")).to(device)

# 计算距离
distance = model(img1, img2) 
print(f"Perceptual Distance: {distance.item()}")

2. 使用单分支模型（轻量级）

默认模型是 CLIP、DINO 和 OpenCLIP 的集成版。如果需要更快的推理速度（约快 3 倍），可以使用单分支模型（如 dino_vitb16），性能略有取舍。

# 加载 DINO 单分支模型
model, preprocess = dreamsim(pretrained=True, dreamsim_type="dino_vitb16", device="cuda")

img1 = preprocess(Image.open("img1_path")).to("cuda")
img2 = preprocess(Image.open("img2_path")).to("cuda")

distance = model(img1, img2)

3. 提取特征向量 (用于图像检索)

在进行图像检索时，建议先提取嵌入向量，再计算余弦距离，以避免重复计算。

import torch.nn.functional as F

# 提取查询图像的特征
query_img = preprocess(Image.open("query.jpg")).to("cuda")
query_embed = model.embed(query_img)

# 提取待检索图像的特征
target_img = preprocess(Image.open("target.jpg")).to("cuda")
target_embed = model.embed(target_img)

# 计算余弦距离 (1 - 余弦相似度)
# 如果 embeddings 已归一化（默认情况），也可使用 L2 距离
dist = (1 - F.cosine_similarity(query_embed, target_embed, dim=-1)).item()
print(f"Distance: {dist}")

4. 作为损失函数使用

DreamSim 可微分，可直接用作深度学习训练中的感知损失函数（类似 LPIPS）。

# 假设 predicted_image 和 reference_image 是需要优化的张量
dist = model(predicted_image, reference_image)
dist.backward()
optimizer.step()