特征提取和预计算特征训练指南

本文档说明如何使用预提取的特征来加速视频训练。

概述

为了提升训练效率，我们可以预先提取：

视频特征（codes）：使用 CosmosVideoTokenizer 将视频编码为离散token
文本特征（embeddings）：使用 T5/UMT5 将文本编码为embedding

训练时直接加载这些预提取的特征，避免每次训练都重新编码。

步骤1：提取特征

使用 extract_features.py 脚本提取特征：

python train/extract_features.py \
    --csv_path /path/to/OpenVid1M_reorganized.csv \
    --video_root_dir /path/to/video_reorg \
    --output_dir /path/to/extracted_features \
    --text_encoder_architecture umt5-base \
    --video_tokenizer_model_id Cosmos-1.0-Tokenizer-DV8x16x16 \
    --num_frames 16 \
    --video_height 480 \
    --video_width 848 \
    --batch_size 4 \
    --num_workers 4

参数说明

--csv_path: OpenVid1M CSV文件路径
--video_root_dir: 视频文件根目录（可选，会自动检测）
--output_dir: 特征保存目录
--text_encoder_architecture: 文本编码器架构（umt5-base/umt5-xxl/t5）
--video_tokenizer_model_id: Cosmos视频tokenizer模型ID
--num_frames, --video_height, --video_width: 视频参数
--batch_size: 批处理大小（根据GPU内存调整）
--num_workers: 数据加载器工作进程数
--max_samples: 最大处理样本数（用于测试，默认处理全部）
--resume_from_index: 从指定索引恢复提取（用于中断后恢复）

输出结构

提取的特征会保存在 output_dir 下，采用三层目录结构以避免单个文件夹下文件过多：

extracted_features/
├── video_codes/          # 视频codes（三层目录结构）
│   ├── 000/             # 第一层：index // 1000000
│   │   ├── 000/         # 第二层：(index // 1000) % 1000
│   │   │   ├── 000/     # 第三层：index % 1000
│   │   │   │   ├── 00000000.npy
│   │   │   │   └── ...
│   │   │   └── ...
│   │   └── ...
│   └── ...
├── text_embeddings/      # 文本embeddings（三层目录结构）
│   ├── 000/
│   │   ├── 000/
│   │   │   ├── 000/
│   │   │   │   ├── 00000000.npy
│   │   │   │   └── ...
│   │   │   └── ...
│   │   └── ...
│   └── ...
└── metadata.json         # 元数据（包含样本信息）

目录结构说明：

对于索引 index，文件路径为：level1/level2/level3/index.npy
- level1 = index // 1000000 (0-999)
- level2 = (index // 1000) % 1000 (0-999)
- level3 = index % 1000 (0-999)
例如：索引 1234567 的文件路径为 001/234/567/1234567.npy

这种结构可以支持最多1,000,000,000个样本，每层最多1000个文件夹，避免单个文件夹下文件过多的问题。

步骤2：使用预提取特征训练

在训练脚本中使用 --use_precomputed_features 和 --features_dir 参数：

python train/train_mei_video.py \
    --use_precomputed_features \
    --features_dir /path/to/extracted_features \
    --text_encoder_architecture umt5-base \
    --video_tokenizer_model_id Cosmos-1.0-Tokenizer-DV8x16x16 \
    --num_frames 16 \
    --video_height 480 \
    --video_width 848 \
    --train_batch_size 8 \
    --learning_rate 3e-4 \
    --max_train_steps 10000 \
    --output_dir ./output \
    --mixed_precision bf16 \
    --gradient_checkpointing \
    --wan_pretrained_path /path/to/wan/weights \
    --wan_backbone_lr_ratio 0.1 \
    --freeze_wan_backbone  # 可选：冻结backbone

关键参数

--use_precomputed_features: 启用预提取特征模式
--features_dir: 预提取特征的目录路径
其他训练参数保持不变

优势

训练速度提升：避免每次训练都重新编码视频和文本
GPU利用率提升：减少CPU-GPU数据传输
内存效率：特征文件比原始视频小得多
可重复性：使用相同的特征确保训练一致性

注意事项

特征一致性：确保提取特征时使用的参数（num_frames, height, width, text_encoder）与训练时一致
存储空间：1M样本的特征大约需要几十GB存储空间
恢复提取：如果提取中断，可以使用 --resume_from_index 参数恢复
验证：训练时仍需要text_encoder进行验证，但不会用于训练数据编码

故障排除

特征提取失败

检查视频文件路径是否正确
检查GPU内存是否足够（可以减小batch_size）
查看日志中的错误信息

训练时找不到特征

确认 --features_dir 路径正确
确认特征文件存在（检查 video_codes/ 和 text_embeddings/ 目录）
检查 metadata.json 文件是否存在

维度不匹配

确保提取和训练时使用相同的视频参数（num_frames, height, width）
确保使用相同的text_encoder架构