VideoGPT：使用 VQ-VAE 和 Transformers 进行视频生成

[论文][项目网站][Colab 教程]
已集成至 Huggingface Spaces，使用 Gradio。查看演示：

我们提出 VideoGPT：一种概念上简洁的架构，用于将基于似然（likelihood-based）的生成建模扩展到自然视频领域。VideoGPT 使用 VQ-VAE（Vector Quantized Variational AutoEncoder，向量量化变分自编码器），通过 3D 卷积和轴向自注意力机制学习原始视频的下采样离散潜在表示。随后，采用类似 GPT 的简单架构，利用时空位置编码对离散潜在变量进行自回归建模。尽管公式简洁且训练容易，我们的架构在 BAIR Robot 数据集上的视频生成效果可与当前最先进的 GAN 模型相媲美，并能从 UCF-101 和 Tumblr GIF 数据集（TGIF）中生成高保真自然图像。我们希望所提出的架构能作为基于 Transformer 的视频生成模型的最小化、可复现实现参考。

方法

安装

请根据你的机器调整 cudatoolkit 版本。

conda install --yes -c conda-forge cudatoolkit=11.0 cudnn
pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html
pip install git+https://github.com/wilson1yan/VideoGPT.git

稀疏注意力（Sparse Attention，可选）

在计算资源受限的情况下，使用稀疏注意力可能更有利。

sudo apt-get install llvm-9-dev
DS_BUILD_SPARSE_ATTN=1 pip install deepspeed

安装 deepspeed 后，可通过在 scripts/train_videogpt.py 中设置标志 --attn_type sparse 来训练稀疏 Transformer。默认支持的稀疏配置是 N 维步幅稀疏布局，但你也可以编写自己的任意稀疏布局。

数据集

默认代码接受符合 videogpt/data.py 中指定格式的 HDF5 文件。例如，可以通过运行以下脚本从 BAIR Robot 数据构建此类数据集：

sh scripts/preprocess/bair/create_bair_dataset.sh datasets/bair

或者，代码也支持如下目录结构的数据集：

video_dataset/
    train/
        class_0/
            video1.mp4
            video2.mp4
            ...
        class_1/
            video1.mp4
            ...
        ...
        class_n/
            ...
    test/
        class_0/
            video1.mp4
            video2.mp4
            ...
        class_1/
            video1.mp4
            ...
        ...
        class_n/
            ...

例如，可通过运行以下脚本从 UCF-101 数据构建此类数据集：

sh scripts/preprocess/ucf101/create_ucf_dataset.sh datasets/ucf101

你可能需要安装 unrar 和 unzip 才能使代码正常工作。

如果你不关心类别，可以省略类别文件夹，直接将数据集结构简化为仅包含视频文件的 train 和 test 目录。

from torchvision.io import read_video
from videogpt import load_vqvae
from videogpt.data import preprocess

video_filename = 'path/to/video_file.mp4'
sequence_length = 16
resolution = 128
device = torch.device('cuda')

vqvae = load_vqvae('kinetics_stride2x4x4')
video = read_video(video_filename, pts_unit='sec')[0]
video = preprocess(video, resolution, sequence_length).unsqueeze(0).to(device)

encodings = vqvae.encode(video)
video_recon = vqvae.decode(encodings)

训练 VQ-VAE

使用 scripts/train_vqvae.py 脚本训练 VQ-VAE。执行 python scripts/train_vqvae.py -h 可查看所有可用训练设置。以下是部分更相关的设置及其默认值：

VQ-VAE 专属设置

• --embedding_dim：码本嵌入的维度数
• --n_codes 2048：码本中的码字数量
• --n_hiddens 240：残差块中的隐藏特征数
• --n_res_layers 4：残差块的数量
• --downsample 4 4 4：编码器在时间 T、高度 H、宽度 W 上的下采样步长

训练设置

• --gpus 2：用于分布式训练的 GPU 数量
• --sync_batchnorm：当使用 >1 个 GPU 时，使用 SyncBatchNorm 替代 BatchNorm3d
• --gradient_clip_val 1：训练时的梯度裁剪阈值
• --batch_size 16：每个 GPU 的批大小
• --num_workers 8：每个 DataLoader 的工作进程数

数据集设置

• --data_path <path>：指向 hdf5 文件或包含 train 和 test 子目录（内含视频子目录）的文件夹路径
• --resolution 128：训练的空间分辨率
• --sequence_length 16：时间分辨率，即视频片段长度

训练 VideoGPT

使用 scripts/train_videogpt.py 脚本训练 VideoGPT 模型以进行采样。执行 python scripts/train_videogpt.py -h 可查看所有可用训练设置。以下是部分更相关的设置及其默认值：

VideoGPT 专属设置

• --vqvae kinetics_stride4x4x4：VQ-VAE 检查点文件路径
• --n_cond_frames 0：条件帧数量。0 表示非帧条件模型
• --class_cond：若启用，则训练类别条件模型
• --hidden_dim 576：Transformer 隐藏特征数
• --heads 4：多头注意力的头数
• --layers 8：Transformer 层数
• --dropout 0.2'：注意力层和位置前馈层后特征的 dropout 概率
• --attn_type full：注意力类型，可选 full 或 sparse。安装稀疏注意力请参阅“安装”章节
• --attn_dropout 0.3：注意力权重矩阵的 dropout 概率

训练设置

• --gpus 4：用于分布式训练的 GPU 数量
• --gradient_clip_val 1：训练时的梯度裁剪阈值
• --batch_size 8：每个 GPU 的批大小
• --num_workers 2：每个 DataLoader 的工作进程数
• --amp_level O1：用于混合精度训练
• --precision 16：用于混合精度训练

数据集设置

• --data_path <path>：指向 hdf5 文件或包含 train 和 test 子目录（内含视频子目录）的文件夹路径
• --resolution 128：训练的空间分辨率
• --sequence_length 16：时间分辨率，即视频片段长度

采样 VideoGPT

VideoGPT 模型可通过 scripts/sample_videogpt.py 脚本进行采样。你可以在训练过程中指定一个检查点（checkpoint）路径。你可能需要安装 ffmpeg：sudo apt-get install ffmpeg

评估

评估主要使用 Frechet Video Distance (FVD，弗雷歇视频距离) 对 BAIR 和 Kinetics 数据集进行衡量，对 UCF-101 数据集则使用 Inception Score（初始分数）。Inception Score 可通过生成样本并使用 TGANv2 仓库 中的代码计算。FVD 可通过运行 python scripts/compute_fvd.py 计算，该脚本使用了从原始代码库移植到 PyTorch 的版本。

复现论文结果

请注意，本仓库主要设计目标是简洁性及便于在此方法基础上扩展。如需完整复现论文中的结果，请使用另一个独立仓库中的代码。但请注意，该仓库的代码整洁度较低。

引用

使用我们的代码时，请考虑引用以下文献：

@misc{yan2021videogpt,
      title={VideoGPT: Video Generation using VQ-VAE and Transformers}, 
      author={Wilson Yan and Yunzhi Zhang and Pieter Abbeel and Aravind Srinivas},
      year={2021},
      eprint={2104.10157},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

VideoGPT 快速上手指南

环境准备

• 操作系统：Linux（推荐 Ubuntu 18.04+）
• GPU 支持：需 NVIDIA 显卡，支持 CUDA 11.0 或兼容版本
• 前置依赖：
  • Python 3.7+
  • Conda（推荐 Miniconda）
  • PyTorch 1.7.1+cu110（如使用其他 CUDA 版本，请调整安装命令）
  • 可选：llvm-9-dev、ffmpeg、unrar、unzip（用于数据预处理和稀疏注意力）

国内用户建议配置清华或阿里云 PyPI 镜像加速 pip 安装：

pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

# 1. 安装 CUDA 工具包（根据你的系统 CUDA 版本调整）
conda install --yes -c conda-forge cudatoolkit=11.0 cudnn

# 2. 安装 PyTorch（CUDA 11.0 版本）
pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html

# 3. 安装 VideoGPT
pip install git+https://github.com/wilson1yan/VideoGPT.git

# （可选）安装稀疏注意力支持
sudo apt-get install llvm-9-dev
DS_BUILD_SPARSE_ATTN=1 pip install deepspeed

基本使用

1. 加载预训练 VQ-VAE 并重建视频片段

from torchvision.io import read_video
from videogpt import load_vqvae
from videogpt.data import preprocess
import torch

video_filename = 'path/to/video_file.mp4'
sequence_length = 16
resolution = 128
device = torch.device('cuda')

vqvae = load_vqvae('kinetics_stride2x4x4')  # 加载预训练模型
video = read_video(video_filename, pts_unit='sec')[0]
video = preprocess(video, resolution, sequence_length).unsqueeze(0).to(device)

encodings = vqvae.encode(video)      # 编码为离散潜变量
video_recon = vqvae.decode(encodings) # 解码重建视频

2. 训练你自己的 VQ-VAE（示例命令）

python scripts/train_vqvae.py \
    --data_path datasets/bair \
    --resolution 128 \
    --sequence_length 16 \
    --batch_size 16 \
    --gpus 2

3. 训练 VideoGPT 模型（需先有 VQ-VAE 权重）

python scripts/train_videogpt.py \
    --vqvae path/to/vqvae_checkpoint.ckpt \
    --data_path datasets/bair \
    --resolution 128 \
    --sequence_length 16 \
    --batch_size 8 \
    --gpus 4

更多训练参数可通过 python scripts/train_videogpt.py -h 查看。