VidTok - 一系列多功能和最先进的视频 tokenizer

一、关于 VidTok

VidTok 是一个尖端的视频标记器系列，在连续和离散标记化方面都表现出色。VidTok结合了对现有方法的几个关键进步：

• ⚡高效的架构。分离的空间和时间采样在不牺牲质量的情况下降低了计算复杂性。
• 🔥高级量化。有限标量量化（FSQ）解决了离散标记化中的训练不稳定性和密码本崩溃。
• 💥增强训练。两阶段策略——低分辨率视频的预训练和高分辨率视频的微调——提高了效率。降低帧速率改善了运动动力学表现。

VidTok在大规模视频数据集上进行了训练，在所有指标（包括PSNR、SSIM、LPIPS和FVD）上都优于以前的模型。

• github : https://github.com/microsoft/VidTok
• paper : https://arxiv.org/pdf/2412.13061
• HuggingFace : https://huggingface.co/microsoft/VidTok

演示：
vidtok_compare.mp4

二、设置

1、克隆此存储库并导航到VidTok文件夹：

git clone https://github.com/microsoft/VidTok
cd VidTok

2、我们提供了一个environment.yaml文件来设置Conda环境。此处提供了Conda的安装说明。

# 1. Prepare conda environment
conda env create -f environment.yaml
# 2. Activate the environment
conda activate vidtok

三、Checkpoints

每个检查点在configs文件夹中都有一个对应的同名configs文件。

└── checkpoints
    ├── vidtok_fsq_causal_41616_262144.ckpt
    ├── vidtok_fsq_causal_488_262144.ckpt
    ├── vidtok_fsq_causal_488_32768.ckpt
    ├── vidtok_fsq_causal_488_4096.ckpt
    ├── vidtok_fsq_noncausal_41616_262144.ckpt
    ├── vidtok_fsq_noncausal_488_262144.ckpt
    ├── vidtok_kl_causal_41616_4chn.ckpt
    ├── vidtok_kl_causal_488_16chn.ckpt
    ├── vidtok_kl_causal_488_8chn.ckpt
    ├── vidtok_kl_causal_488_4chn.ckpt
    ├── vidtok_kl_noncausal_41616_4chn.ckpt
    └── vidtok_kl_noncausal_488_4chn.ckpt

每个检查点在configs文件夹中都有一个对应的同名configs文件。

• VCR表示视频压缩比TxHxW。
• 上表显示了在MCL_JCL数据集中的30个测试视频上评估的模型性能，样本fps为30。因果模型的输入大小为17x256x256，非因果模型的输入大小为16x256x256。

四、训练

1、数据准备

1、将所有培训视频DATA_DIR：

└── DATA_DIR
    ├── subset1
    │   ├── videoname11.mp4
    │   └── videoname12.mp4
    ├── subset2
    │   ├── videoname21.mp4
    │   ├── videoname22.mp4
    │   └── subsubset1
    │       ├── videoname211.mp4
    │       └── videoname212.mp4
    └── ...

2、准备一个.csv元文件来记录这些视频相对于DATA_DIR的相对路径，例如：

videos
subset1/videoname11.mp4
subset2/videoname21.mp4
subset2/subsubset1/videoname211.mp4

验证数据也按照上述步骤准备。

2、微调自定义数据

1、根据数据准备准备您自己的训练和验证数据。

2、根据需要从configs文件夹中选择合适的CONFIG文件，并修改以下参数：

指定ckpt_path参数以使用预训练的检查点参数初始化模型：

model:
  params:
    ckpt_path: PATH_TO_CHECKPOINT  # train from existing checkpoint

指定data部分以使用您自己的训练和验证数据：

  train:
      target: vidtok.data.vidtok.VidTokDataset
      params:
        data_dir: DATA_DIR_1  # DATA_DIR for training data
        meta_path: META_PATH_1  # path to the .csv meta file of training data
        video_params:
          input_height: INPUT_HEIGHT_1
          input_width: INPUT_WIDTH_1
          sample_num_frames: NUM_FRAMES_1  # typically set to 17 for causal models and 16 for non-causal models
          sample_fps: SAMPLE_FPS_1  # sample fps for training data
  validation:
      target: vidtok.data.vidtok.VidTokDataset
      params:
        data_dir: DATA_DIR_2  # DATA_DIR for validation data
        meta_path: META_PATH_2  # path to the .csv meta file of validation data
        video_params: 
          input_height: INPUT_HEIGHT_2
          input_width: INPUT_WIDTH_2
          sample_num_frames: NUM_FRAMES_2  # typically set to 17 for causal models and 16 for non-causal models
          sample_fps: SAMPLE_FPS_2  # sample fps for validation data
        start_index: 0  # fixed value to ensure the same sampled data

将fix_encoder和fix_decoder设置为False以启用全模型微调：

model:
  params:
    encoder_config:
      params:
        fix_encoder: false
        fix_decoder: false

其他超参数根据您的需要。

3、运行以下命令开始训练：

python main.py -b CONFIG --logdir LOGDIR

# You can also use `torchrun` to start the training code.

训练日志和检查点保存在LOGDIR。

建议使用权重&偏差作为数据可视化工具（默认为TensorBoard），使用wandb login先登录，然后运行：

python main.py -b CONFIG --logdir LOGDIR --wandb --wandb_entity ENTITY --wandb_project PROJECT

3、从零开始的训练

五、推理

1、简单使用

我们提供以下示例以快速使用我们的模型。它适用于连续和离散标记化以及因果和非因果模型。只需提供配置文件cfg_path和检查点文件ckpt_path的路径，并相应地将is_causal设置为True或False。

import torch
from scripts.inference_evaluate import load_model_from_config

cfg_path = "configs/vidtok_kl_causal_488_4chn.yaml"
ckpt_path = "checkpoints/vidtok_kl_causal_488_4chn.ckpt"
is_causal = True

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")    
# load pre-trained model
model = load_model_from_config(cfg_path, ckpt_path)
model.to(device).eval()
# random input
num_frames = 17 if is_causal else 16
x_input = (torch.rand(1, 3, num_frames, 256, 256) * 2 - 1).to(device)  # [B, C, T, H, W], range -1~1
# model forward
_, x_recon, _ = model(x_input)
assert x_input.shape == x_recon.shape

2、重建输入视频

python scripts/inference_reconstruct.py --config CONFIG --ckpt CKPT --input_video_path VIDEO_PATH --num_frames_per_batch NUM_FRAMES_PER_BATCH --input_height 256 --input_width 256 --sample_fps 30 --output_video_dir OUTPUT_DIR

• 指定测试视频路径的VIDEO_PATH。我们在assets/example.mp4中提供了一个示例视频。
• 因果模型NUM_FRAMES_PER_BATCH为17，非因果模型为16。
• 重建的视频保存在OUTPUT_DIR中。
• 对于因果模型，您可以选择在命令行中添加--pad_gen_frames，这可能会提高重建视频的流畅度。

3、绩效评估

我们还提供了一个手稿scripts/inference_evaluate.py来评估PSNR、SSIM和LPIPS中的视频重建性能。

1. 把你所有的测试视频DATA_DIR。
2. 运行以下命令，将测试DATA_DIR下的所有.mp4视频：

python scripts/inference_evaluate.py --config CONFIG --ckpt CKPT --data_dir DATA_DIR --num_frames_per_batch NUM_FRAMES_PER_BATCH --input_height 256 --input_width 256 --sample_fps 30

（可选）如果只想测试DATA_DIR下的某些视频，则需要准备一个.csv元文件来指示要测试的视频文件（参考数据准备）。并在上面的命令中添加--meta_path META_PATH来指定.csv元文件的路径。

六、其它

预期用途

我们正在与研究界分享我们的模型，以促进该领域的进一步研究：

• 训练您自己的视频标记器用于研究目的。
• 具有各种压缩率的视频标记化。

范围外用途

我们的模型不是专门为所有下游目的设计或评估的。开发人员在选择用例时应考虑视频标记器的常见限制（例如，域外数据的性能下降），并在特定下游用例中使用之前评估和减轻隐私、安全性和公平性，特别是对于高风险场景。

开发人员应了解并遵守与其用例相关的适用法律或法规（包括隐私、贸易合规法等）。

风险和限制

需要注意的此模型的一些限制包括：

• VidTok可能会丢失有关重建内容的详细信息。
• VidTok继承其训练数据的任何偏差、错误或遗漏。
• VidTok是为研究和实验目的而开发的。在考虑将其应用于商业或现实世界场景之前，需要进一步的测试和验证。

2025-01-06(一)