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

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

Meta AI 研究团队，FAIR

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

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

V-JEPA 模型通过被动观看来自 VideoMix2M 数据集的视频像素进行训练，并产生通用的视觉表示，这些表示在下游视频和图像任务（downstream tasks）中表现良好，而无需调整模型的参数；例如，使用冻结的主干网络（backbone）和一个轻量级的特定任务注意力探针（attentive probe）。

方法

V-JEPA 预训练仅基于无监督特征预测目标（unsupervised feature prediction objective），不利用预训练的图像编码器（pretrained image encoders）、文本、负样本（negative examples）、人工标注或像素级重建（pixel-level reconstruction）。

     

可视化

与具有像素解码器（pixel decoder）的生成式方法（generative methods）不同，V-JEPA 拥有一个在潜在空间（latent space）中进行预测的预测器（predictor）。 我们训练一个条件扩散模型（conditional diffusion model）将 V-JEPA 特征空间（feature-space）的预测解码为可解释的像素；在此过程中，预训练的 V-JEPA 编码器和预测器网络保持冻结。 该解码器仅接收针对视频中缺失区域预测的表示，无法访问未掩码（unmasked）的视频区域。

V-JEPA 的特征预测确实具有依据，并与视频的未掩码区域表现出时空一致性（spatio-temporal consistency）。

模型库

预训练模型

模型	补丁大小	分辨率	迭代次数	批次大小	数据	下载
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	注意力探针检查点 (attentive probe checkpoint)	配置文件
ViT-H/16	224x224	75.9	注意力探针检查点	配置文件
ViT-H/16	384x384	77.4	注意力探针检查点	配置文件

Places205 注意力探针 (attentive probes)

模型	分辨率	准确率	下载
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                       # the only place where training loops are allowed
│   ├── vjepa                 #   Video JEPA pre-training
│   ├── main_distributed.py   #   entrypoint for launching app on slurm cluster
│   └── main.py               #   entrypoint for launching app locally on your machine for debugging
├── evals                     # the only place where evaluation of 'apps' are allowed
│   ├── image_classification  #   training an attentive probe for image classification with frozen backbone
│   ├── video_classification  #   training an attentive probe for video classification with frozen backbone
│   ├── main_distributed.py   #   entrypoint for launching distributed evaluations on slurm cluster
│   └── main.py               #   entrypoint for launching evaluations locally on your machine for debugging
├── src                       # the package
│   ├── datasets              #   datasets, data loaders, ...
│   ├── models                #   model definitions
│   ├── masks                 #   mask collators, masking utilities, ...
│   └── utils                 #   shared utilities
└── configs                   # the only place where config files are allowed (specify experiment params for app/eval runs)
    ├── evals                 #   configs for launching vjepa frozen evaluations
    └── pretrain              #   configs for launching vjepa pretraining

数据准备

视频数据集

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

/absolute_file_path.[mp4, webvid, etc.] $integer_class_label
/absolute_file_path.[mp4, webvid, etc.] $integer_class_label
/absolute_file_path.[mp4, webvid, etc.] $integer_class_label
...

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

图像数据集

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

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

.
└── /my_image_datasets/                # where we store image datasets
    ├── places205/121517/pytorch/      #   Places205
    │   └── [...]
    ├── iNaturalist-2021/110421/       #   iNaturalist21
    │   └── [...]
    ├── [...]                          #   Other Image Datasets
    │   └── [...]
    └── imagenet_full_size/061417/     #   ImageNet1k
        └── train
        │   ├── $class_1
        │   │    ├── xxx.[png, jpeg, etc.]
        │   │    ├── [...]
        │   │    └── xxz.[png, jpeg, etc.]
        │   ├── [...]
        │   └── $class_n
        │       ├── abc.[png, jpeg, etc.]
        │       ├── [...]
        │       └── abz.[png, jpeg, etc.]
        └── val
            ├── $class_1
            │    ├── xxx.[png, jpeg, etc.]
            │    ├── [...]
            │    └── xxz.[png, jpeg, etc.]
            ├── [...]
            └── $class_n
                ├── abc.[png, jpeg, etc.]
                ├── [...]
                └── abz.[png, jpeg, etc.]

启动 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}
}

V-JEPA 快速上手指南

V-JEPA (Video Joint Embedding Predictive Architecture) 是由 Meta AI Research (FAIR) 推出的视频自监督学习框架。它通过被动观看视频像素进行无监督特征预测，无需图像编码器、文本或负样本即可生成通用的视觉表示。

1. 环境准备

本项目基于 PyTorch 开发，建议配置如下：

• 操作系统: Linux / macOS / Windows (推荐 Linux)
• Python: 3.8+
• 深度学习框架: PyTorch (请根据官方文档选择对应 CUDA 版本)
• 硬件: 支持 CUDA 的 GPU (多卡训练需配置分布式环境)

2. 安装步骤

克隆代码库

git clone https://github.com/facebookresearch/jepa.git
cd jepa

安装依赖

确保已安装 PyTorch，并安装项目所需的其他依赖（通常位于根目录）：

pip install -r requirements.txt

(注：若未找到 requirements.txt，请参考官方文档安装基础依赖)

数据准备

V-JEPA 支持多种标准视频格式。将视频数据集转换为兼容格式，创建一个 .csv 文件，并在配置文件中指定路径：

/absolute/path/to/video.mp4,label

(具体字段格式请参考原始 README 中的完整说明)

3. 基本使用

所有实验参数均通过配置文件管理，而非命令行参数。

修改配置

在启动前，必须更新配置文件中的路径，指向您的日志保存目录、检查点目录及训练数据位置。

• 预训练配置：configs/pretrain/
• 评估配置：configs/evals/

本地调试运行

使用 main.py 在本地机器上启动应用：

# 预训练入口
python app/main.py

# 评估入口
python evals/main.py

分布式集群运行

在 Slurm 集群上使用分布式脚本启动：

# 预训练分布式启动
python app/main_distributed.py

# 评估分布式启动
python evals/main_distributed.py

4. 预训练模型与探针

项目提供了多个预训练的 ViT 模型及下游任务探针（Attentive Probes），可直接下载用于推理或微调。

模型	分辨率	任务	下载链接
ViT-L/16	224x224	K400 Probe	checkpoint
ViT-H/16	224x224	K400 Probe	checkpoint
ViT-L/16	224x224	ImageNet Probe	checkpoint
ViT-H/16	384x384	SSv2 Probe	checkpoint

更多模型及详细配置请查看 MODEL ZOO 表格。