V-JEPA：视频联合嵌入预测架构

PyTorch官方代码库，用于视频联合嵌入预测架构V-JEPA，这是一种从视频中自监督学习视觉表示的方法。

Meta AI研究院，FAIR

Adrien Bardes、Quentin Garrido、Jean Ponce、Xinlei Chen、Michael Rabbat、Yann LeCun、Mahmoud Assran*、Nicolas Ballas*

[博客] [论文] [Yannic Kilcher的视频]

V-JEPA模型通过被动观看来自VideoMix2M数据集的视频像素进行训练，产生多功能的视觉表示，这些表示在下游视频和图像任务上表现出色，无需调整模型参数；例如，使用冻结的骨干网络和仅有的轻量级任务特定注意力探针。

方法

V-JEPA预训练仅基于无监督特征预测目标，不使用预训练的图像编码器、文本、负样本、人工标注或像素级重建。

可视化

与具有像素解码器的生成方法不同，V-JEPA有一个在潜在空间中进行预测的预测器。我们训练了一个条件扩散模型，将V-JEPA特征空间预测解码为可解释的像素；在此过程中，预训练的V-JEPA编码器和预测器网络保持冻结状态。解码器仅输入视频缺失区域的预测表示，无法访问视频的未遮蔽区域。

V-JEPA的特征预测确实是有根据的，并且与视频未遮蔽区域表现出时空一致性。

模型库

预训练模型

模型	patch大小	分辨率	迭代次数	批量大小	数据	下载
ViT-L	2x16x16	224x224	90K	3072	VideoMix2M	检查点	配置
ViT-H	2x16x16	224x224	90K	3072	VideoMix2M	检查点	配置
ViT-H	2x16x16	384x384	90K	2400	VideoMix2M	检查点	配置

K400注意力探针

模型	分辨率	准确率 (16x8x3)	下载
ViT-L/16	224x224	80.8	注意力探针检查点	配置
ViT-H/16	224x224	82.0	注意力探针检查点	配置
ViT-H/16	384x384	81.9	注意力探针检查点	配置

SSv2注意力探针

模型	分辨率	准确率 (16x2x3)	下载
ViT-L/16	224x224	69.5	注意力探针检查点	配置
ViT-H/16	224x224	71.4	注意力探针检查点	配置
ViT-H/16	384x384	72.2	注意力探针检查点	配置

ImageNet1K注意力探针

模型	分辨率	准确率	下载
ViT-L/16	224x224	74.8	注意力探针检查点	配置
ViT-H/16	224x224	75.9	注意力探针检查点	配置
ViT-H/16	384x384	77.4	注意力探针检查点	配置

Places205 注意力探针

模型	分辨率	准确率	下载
ViT-L/16	224x224	60.3	注意力探针检查点	配置
ViT-H/16	224x224	61.7	注意力探针检查点	配置
ViT-H/16	384x384	62.8	注意力探针检查点	配置

iNat21 注意力探针

模型	分辨率	准确率	下载
ViT-L/16	224x224	67.8	注意力探针检查点	配置
ViT-H/16	224x224	67.9	注意力探针检查点	配置
ViT-H/16	384x384	72.6	注意力探针检查点	配置

代码结构

配置文件: 所有实验参数都在配置文件中指定（而不是命令行参数）。查看 configs/ 目录以获取示例配置文件。注意，在启动实验之前，您必须更新配置文件中的路径，指向您自己的目录，指示在哪里保存日志和检查点，以及在哪里找到训练数据。

.
├── app                       # 唯一允许训练循环的地方
│   ├── vjepa                 #   视频 JEPA 预训练
│   ├── main_distributed.py   #   在 slurm 集群上启动应用程序的入口点
│   └── main.py               #   在本地机器上启动应用程序以进行调试的入口点
├── evals                     # 唯一允许评估"应用程序"的地方
│   ├── image_classification  #   使用冻结的主干网络训练注意力探针进行图像分类
│   ├── video_classification  #   使用冻结的主干网络训练注意力探针进行视频分类
│   ├── main_distributed.py   #   在 slurm 集群上启动分布式评估的入口点
│   └── main.py               #   在本地机器上启动评估以进行调试的入口点
├── src                       # 包
│   ├── datasets              #   数据集、数据加载器等
│   ├── models                #   模型定义
│   ├── masks                 #   掩码整理器、掩码实用程序等
│   └── utils                 #   共享实用程序
└── configs                   # 唯一允许配置文件的地方（指定应用程序/评估运行的实验参数）
    ├── evals                 #   用于启动 vjepa 冻结评估的配置
    └── pretrain              #   用于启动 vjepa 预训练的配置

数据准备

视频数据集

V-JEPA 预训练和评估可以使用许多标准视频格式。要使视频数据集与 V-JEPA 代码库兼容，您只需创建一个具有以下格式的 .csv 文件，然后在配置中指定此 CSV 文件的路径。

/绝对文件路径.[mp4, webvid, 等] $整数类别标签
/绝对文件路径.[mp4, webvid, 等] $整数类别标签
/绝对文件路径.[mp4, webvid, 等] $整数类别标签
...

由于 V-JEPA 完全是无监督的，预训练代码将忽略 CSV 文件中的 $整数类别标签。因此，您可以在这一列中放入随机值。但是，如果您希望在视频数据集上运行有监督的视频分类评估，则必须将 $整数类别标签替换为每个视频的真实标签。

图像数据集

我们在图像分类评估中使用标准的 PyTorch ImageFolder 类。因此，要为图像分类评估设置图像数据集，首先创建一个目录来存储您的图像数据集 $存储图像数据集的目录。接下来，将您的图像数据集下载到此目录中，格式与 PyTorch ImageFolder 兼容。

例如，假设我们有一个名为 my_image_datasets 的目录。然后我们将图像数据集下载到这个目录中，最终得到以下文件树结构：

.
└── /my_image_datasets/                # 存储图像数据集的位置
    ├── places205/121517/pytorch/      #   Places205
    │   └── [...]
    ├── iNaturalist-2021/110421/       #   iNaturalist21
    │   └── [...]
    ├── [...]                          #   其他图像数据集
    │   └── [...]
    └── imagenet_full_size/061417/     #   ImageNet1k
        └── train
        │   ├── $class_1
        │   │    ├── xxx.[png, jpeg, 等]
        │   │    ├── [...]
        │   │    └── xxz.[png, jpeg, 等]
        │   ├── [...]
        │   └── $class_n
        │       ├── abc.[png, jpeg, 等]
        │       ├── [...]
        │       └── abz.[png, jpeg, 等]
        └── val
            ├── $class_1
            │    ├── xxx.[png, jpeg, 等]
            │    ├── [...]
            │    └── xxz.[png, jpeg, 等]
            ├── [...]
            └── $class_n
                ├── abc.[png, jpeg, 等]
                ├── [...]
                └── abz.[png, jpeg, 等]

启动 V-JEPA 预训练

本地训练

如果您希望在启动分布式训练运行之前调试代码或设置，我们提供了在多GPU（或单GPU）机器上本地运行预训练脚本的功能，但复现我们的结果需要启动分布式训练。

单机实现从app/main.py开始，它解析实验配置文件并在多GPU（或单GPU）机器上本地运行预训练。例如，要在本地机器上使用配置configs/pretrain/vitl16.yaml在GPU "0"、"1"和"2"上运行V-JEPA预训练，请输入以下命令：

python -m app.main \
  --fname configs/pretrain/vitl16.yaml \
  --devices cuda:0 cuda:1 cuda:2

分布式训练

要启动分布式训练运行，实现从app/main_distributed.py开始，除了解析配置文件外，还允许指定分布式训练的详细信息。对于分布式训练，我们使用流行的开源工具submitit，并为SLURM集群提供示例。

例如，要使用配置configs/pretrain/vitl16.yaml启动分布式预训练实验，请输入以下命令：

python -m app.main_distributed \
  --fname configs/pretrain/vitl16.yaml \
  --folder $path_to_save_stderr_and_stdout \
  --partition $slurm_partition

启动评估

本地训练

如果您希望在启动分布式训练运行之前调试评估代码或设置，我们提供了在多GPU（或单GPU）机器上本地运行评估脚本的功能，但复现完整评估需要启动分布式训练。单机实现从eval/main.py开始，它解析实验配置文件并在多GPU（或单GPU）机器上本地运行评估。

例如，要在本地机器上使用配置configs/eval/vitl16_in1k.yaml在GPU "0"、"1"和"2"上运行ImageNet图像分类，请输入以下命令：

python -m evals.main \
  --fname configs/eval/vitl16_in1k.yaml \
  --devices cuda:0 cuda:1 cuda:2

分布式训练

要启动分布式评估运行，实现从eval/main_distributed.py开始，除了解析配置文件外，还允许指定分布式训练的详细信息。对于分布式训练，我们使用流行的开源工具submitit，并为SLURM集群提供示例。

例如，要使用配置configs/eval/vitl16_in1k.yaml启动分布式ImageNet图像分类实验，请输入以下命令：

python -m evals.main_distributed \
  --fname configs/eval/vitl16_in1k.yaml \
  --folder $path_to_save_stderr_and_stdout \
  --partition $slurm_partition

同样，要使用配置configs/eval/vitl16_k400.yaml启动分布式K400视频分类实验，请输入以下命令：

python -m evals.main_distributed \
  --fname configs/eval/vitl16_k400.yaml \
  --folder $path_to_save_stderr_and_stdout \
  --partition $slurm_partition

设置

运行：

conda create -n jepa python=3.9 pip
conda activate jepa
python setup.py install

许可证

有关此代码发布的许可证详细信息，请参见LICENSE文件。

引用

如果您在研究中发现此存储库有用，请考虑给予星标:star:并引用

@article{bardes2024revisiting,
  title={Revisiting Feature Prediction for Learning Visual Representations from Video},
  author={Bardes, Adrien and Garrido, Quentin and Ponce, Jean and Rabbat, Michael, and LeCun, Yann and Assran, Mahmoud and Ballas, Nicolas},
  journal={arXiv:2404.08471},
  year={2024}
}