⚡ EB-JEPA

基于能量的联合嵌入预测架构

Meta AI Research, FAIR

Basile Terver, Randall Balestriero, Megi Dervishi, David Fan, Quentin Garrido, Tushar Nagarajan,
Koustuv Sinha, Wancong Zhang, Mike Rabbat, Yann LeCun, Amir Bar

一个开源库及教程，用于通过联合嵌入预测架构学习表示，以支持预测和规划。

每个示例（几乎）都是自包含的，在单张 GPU 卡上训练最多只需几小时。

📚 示例

图像 JEPA

基于 CIFAR-10 数据集中的无标签图像进行自监督表示学习，并在分类任务上进行评估。

视频 JEPA

预测序列中下一帧的表示。

AC 视频 JEPA

用于 Two Rooms 环境中的世界建模与规划的 JEPA。

规划示例	任务描述
成功规划示例	从初始状态到达目标状态

🚀 安装

我们使用 uv 进行包管理。

# 安装依赖
uv sync
# 选项 1：激活虚拟环境
source .venv/bin/activate
python main.py
# 选项 2：直接使用 uv 运行
uv run python main.py

如果您需要特定于 conda 的包，可以使用 Conda + uv。

# 创建 Python 3.12 的 conda 环境
conda create -n eb_jepa python=3.12 -y
conda activate eb_jepa
# 以可编辑模式安装包，并包含开发依赖（pytest、black、isort）
uv pip install -e . --group dev

将以下内容添加到您的 ~/.bashrc 文件中，以实现持久化配置。

# 对于 SLURM 作业查找数据集是必需的
export EBJEPA_DSETS=/path/to/eb_jepa/datasets
# 可选：检查点和日志目录
export EBJEPA_CKPTS=/path/to/checkpoints

🏋️ 训练

快速入门

# 本地训练
python -m examples.{image_jepa,video_jepa,ac_video_jepa}.main

我们的默认配置针对 H100 GPU 进行了优化。对于较旧的 GPU（如 A100 或 V100），您可能需要减小批量大小以适应显存。

📂 文件夹结构

所有实验都采用统一的文件夹结构：

checkpoints/
└── {example_name}/
    ├── dev_2026-01-16_00-10/                 # 单次/本地运行（dev_ 前缀）
    │   └── {exp_name}_seed1/
    │
    ├── sweep_2026-01-16_00-10/         # 自动命名的三种子试验
    │   ├── {exp_name}_seed1/
    │   ├── {exp_name}_seed1000/
    │   └── {exp_name}_seed10000/
    │
    └── sweep_my_experiment/            # 自定义名称的试验
        └── ...

{exp_name} 编码了关键超参数，以避免文件夹冲突，例如：

• image_jepa: resnet_vicreg_proj_bs256_ep300_ph2048_po2048_std1.0_cov80.0
• video_jepa: resnet_bs64_lr0.001_std10.0_cov100.0
• ac_video_jepa: impala_cov8_std16_simt12_idm1

# 推荐使用 3 种子并结合 wandb 平均（建议）
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml

# 自定义试验名称
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --sweep my_experiment

# 单个作业
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --single

# 完整超参数扫描
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --full-sweep

# 使用 wandb 试验界面进行超参数分析
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --use-wandb-sweep

# 示例：examples/image_jepa/cfgs/default.yaml
sweep:
  param_grid:
    loss.cov_coeff: [0.1, 1.0, 10.0, 100.0]
    loss.std_coeff: [1.0, 10.0]
    meta.seed: [1, 1000, 10000]

🧪 运行测试用例

添加到 eb_jepa 的库必须拥有自己的测试用例（位于 /tests/）。运行测试的方法如下：

使用 uv 同步安装

uv run pytest tests/

使用 conda + uv 安装（不会创建 .venv）

pytest tests/

## 👩‍💻 开发

在贡献代码之前，请使用以下工具格式化您的代码：

```bash
# 移除未使用的导入
autoflake --remove-all-unused-imports -r --in-place .
# 排序导入
python -m isort eb_jepa examples tests
# 格式化代码
python -m black eb_jepa examples tests

📚 引用 EB-JEPA

如果您觉得本仓库有用，请考虑给个 ⭐ 并引用如下文献：

@misc{terver2026lightweightlibraryenergybasedjointembedding,
      title={一种用于基于能量的联合嵌入预测架构的轻量级库},
      author={Basile Terver 和 Randall Balestriero、Megi Dervishi、David Fan、Quentin Garrido、Tushar Nagarajan、Koustuv Sinha、Wancong Zhang、Mike Rabbat、Yann LeCun 和 Amir Bar},
      year={2026},
      eprint={2602.03604},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2602.03604},
}

📄 许可证

EB-JEPA 采用 Apache 许可证。详情请参阅 LICENSE。

EB-JEPA 快速上手指南

EB-JEPA 是由 Meta AI Research (FAIR) 开源的库，用于通过基于能量的联合嵌入预测架构（Energy-Based Joint-Embedding Predictive Architectures）学习表示，适用于预测与规划任务。本指南帮助中国开发者快速搭建环境并运行示例。

环境准备

• 操作系统: Linux (推荐 Ubuntu)
• Python 版本: 3.12
• 硬件要求:
  • 单张 GPU 即可运行示例（默认配置针对 H100 优化）。
  • 若使用较旧显卡（如 A100, V100），可能需要减小 batch size 以适应显存。
• 前置工具:
  • 推荐使用 uv 进行包管理（比 pip/conda 更快）。
  • 可选：Conda（用于管理 Python 环境）。

安装步骤

方案一：使用 uv（推荐）

这是最简洁的安装方式，自动处理虚拟环境和依赖。

# 克隆仓库
git clone https://github.com/facebookresearch/eb_jepa.git
cd eb_jepa

# 同步安装依赖（自动创建 .venv）
uv sync

# 激活虚拟环境
source .venv/bin/activate

# 验证安装（可选）
python main.py --help

方案二：Conda + uv

如果你习惯使用 Conda 管理基础环境：

# 创建并激活 Conda 环境 (Python 3.12)
conda create -n eb_jepa python=3.12 -y
conda activate eb_jepa

# 使用 uv 以可编辑模式安装包及开发依赖
uv pip install -e . --group dev

配置环境变量（可选）

为了在 SLURM 集群或特定路径下运行，建议将以下配置添加到 ~/.bashrc：

# 数据集路径（必填，如需加载本地数据）
export EBJEPA_DSETS=/path/to/eb_jepa/datasets
# 检查点和日志保存路径（可选）
export EBJEPA_CKPTS=/path/to/checkpoints

基本使用

EB-JEPA 提供了三个主要示例：图像 JEPA、视频 JEPA 和 智能体规划 (AC Video JEPA)。每个示例几乎是独立的，单卡训练仅需数小时。

1. 快速开始训练

使用以下命令在本地启动训练（默认使用配置文件中的超参数）：

# 运行图像 JEPA 示例 (CIFAR-10 自监督学习)
python -m examples.image_jepa.main

# 运行视频 JEPA 示例 (预测序列中的下一帧表示)
python -m examples.video_jepa.main

# 运行智能体规划示例 (Two Rooms 环境建模与规划)
python -m examples.ac_video_jepa.main

注意: 如果显存不足，请修改对应的 .yaml 配置文件，减小 batch_size。

2. 查看实验结果

训练完成后，检查点和日志将按以下结构保存：

checkpoints/
└── {example_name}/
    └── dev_YYYY-MM-DD_HH-MM/      # 单次运行
        └── {exp_name}_seed1/      # 包含模型权重和日志

实验文件夹名称 {exp_name} 自动编码了关键超参数（如模型类型、batch size、损失系数等），以避免冲突。

3. 高级用法：超参数扫描 (Sweep)

若需进行多种子实验或超参数搜索，可使用内置的启动器（支持 SLURM 或本地多进程）：

# 启动 3 个不同种子的实验（推荐用于复现性验证）
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml

# 启动完整的超参数网格搜索 (需在 yaml 中配置 sweep.param_grid)
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --full-sweep

# 结合 Wandb 进行可视化分析
python -m examples.launch_sbatch --example image_jepa --fname examples/image_jepa/cfgs/default.yaml --use-wandb-sweep

4. 运行测试

确保环境配置正确，可运行单元测试：

# 如果使用 uv sync 安装
uv run pytest tests/

# 如果使用 conda + uv 安装
pytest tests/