拓冰建站拓冰建站
首页 / 资讯中心 / 正文

Open-Sora 视频自编码器(Video DC-AE)实战指南:从压缩原理到训练与推理配置全解析

Open-Sora 视频自编码器Video DC-AE实战指南从压缩原理到训练与推理配置全解析【免费下载链接】Open-SoraOpen-Sora: Democratizing Efficient Video Production for All项目地址: https://gitcode.com/GitHub_Trending/op/Open-Sora导读本文围绕 Open-Sora 仓库中的 docs/ae.md 文档展开系统讲解项目自研的视频自编码器Video DC-AE基于 DC-AE 架构、面向视频的深度压缩自编码器的设计动机、训练与推理完整流程以及训练/推理配置文件的逐项解读。读完本文你将掌握Video DC-AE 相比 HunyuanVideo VAE 在时空压缩比上的优势与 token 数量削减原理、如何用 8 卡启动从零训练与带判别器的两阶段训练、如何用官方或自训练权重做视频重建推理以及 spatial/temporal tiling分块编码机制的核心参数与调优方法。一、Video DC-AE 是什么压缩比与 token 优势1.1 设计动机大幅提升 AE 压缩比在文生视频 / 图生视频的扩散模型管线中视频自编码器AE负责把高分辨率、长帧数的原始视频压缩为紧凑的隐空间表示latent扩散模型在该隐空间上进行去噪最后由解码器还原为视频。隐空间 token 的多少直接决定了扩散模型训练与推理的开销。受 SANAarxiv 2410.10629启发Open-Sora 的目标是大幅提升 AE 的压缩比。为此项目提出了基于 DC-AEMIT Han Lab 的高效视觉自编码器架构的视频自编码器命名为Video DC-AE其压缩能力为时间维度 4x 压缩空间维度 32x32 压缩即 H 和 W 各 32 倍。作为对照Tencent HunyuanVideo 的 VAE 压缩比为4x8x8时间 4 倍、空间各 8 倍。Video DC-AE 在空间压缩比上显著更高因此在相同 patch size 的前提下扩散模型侧的 token 长度整体减少16x从而同时大幅提升训练与推理速度。说明token 削减 16x 的推导基于假设相同 patch sizes这一前提原文明确注明了该假设条件。1.2 模型命名与源码对应模型名dc-ae-f32t4c128直接编码了核心超参数f32空间压缩 32x、t4时间压缩 4x、c128隐空间通道数 128。该配置在 opensora/models/dc_ae/models/dc_ae.py 的dc_ae_f32工厂函数中硬编码为结构化配置time_compression_ratio4 spatial_compression_ratio32 encoder.block_type[ResBlock,ResBlock,ResBlock,EViTS5_GLU,EViTS5_GLU,EViTS5_GLU] encoder.width_list[128,256,512,512,1024,1024] encoder.depth_list[2,2,2,3,3,3] encoder.normrms3d encoder.is_videoTrue decoder.width_list[128,256,512,512,1024,1024] decoder.depth_list[3,3,3,3,3,3] encoder.temporal_downsample[False,False,False,True,True,False] decoder.temporal_upsample[False,False,False,True,True,False] latent_channels128从上述源码结构可以看出编码器/解码器均为 6 级残差阶段前 3 级使用ResBlock RMS 3D 归一化后 3 级替换为带GLUMBConv的高效EViTS5_GLU块时间下采样只发生在第 4、5 两级对应temporal_downsample[False,False,False,True,True,False]配合InterpolateConv上采样实现 4x 时间压缩。值得注意的另一点是Video DC-AE没有变分variational组件因此训练损失中不包含 KL 项详见后文损失配置。二、数据准备Data Preparation训练与推理共用的数据集准备流程请参考仓库内的 docs/train.md#prepare-dataset 一节其中包含完整的 pexels 45k 数据集下载与解压命令、三个 csv 文件pexels_45k.csv/pexels_45k_necessary.csv/pexels_45k_score.csv的用途说明以及用scripts/cnv/meta.py把原始 csv 处理成训练所需格式的方法。如果使用自定义数据集只需在训练或推理命令中追加参数--dataset.data_path your_data_path自定义数据集 csv 至少需要包含以下列对应 docs/train.md 中的说明path,text,num_frames,height,width,aspect_ratio,resolution,fps仓库自带的训练配置 configs/vae/train/video_dc_ae.py 默认使用datasets/pexels_45k_necessary.csv作为数据源并将fps_max限制为 24。三、模型训练两阶段完整流程3.1 训练规模与总体策略官方训练设定为在 8 张 GPU 上从零开始训练约 3 周。整个训练采用两阶段策略第一阶段只用重建损失 感知损失训练 Video DC-AE 本体无判别器第二阶段模型基本收敛后引入判别器从 checkpoint 继续训练加入生成对抗损失提升重建质量。3.2 第一阶段纯重建训练启动 8 卡训练命令torchrun --nproc_per_node 8 scripts/vae/train.py configs/vae/train/video_dc_ae.py对应的训练入口是 scripts/vae/train.py它基于 ColossalAI 的Booster构建分布式训练环境。基线训练配置 configs/vae/train/video_dc_ae.py 中值得关注的默认项包括配置项默认值说明model.typedc_ae模型注册类型对应 DC-AE 实现model.model_namedc-ae-f32t4c128f32空间 32xt4时间 4xc128隐通道 128model.from_scratchTrue从零初始化不加载预训练权重optim.clsHybridAdamAdam 优化器lr5e-5betas(0.9, 0.98)无 weight decaylr_scheduler.warmup_steps0无 warmupmixed_strategymixed_video_image视频/图像混合训练策略mixed_image_ratio0.2图像样本占比 1:4dtypebf16bfloat16 混合精度pluginzero2ZeRO-2 分布式插件grad_clip1.0梯度裁剪grad_checkpointFalse基线配置下不开启梯度检查点epochs100训练轮数ckpt_every3000每 3000 步保存一次 checkpointema_decay0.99EMA 模型衰减系数pin_memory_cache_pre_alloc_numels[50 * 1024 * 1024] * num_workers * prefetch_factor预分配 pinned memory 缓存3.3 第二阶段加入判别器继续训练当模型接近收敛后使用带判别器的配置 configs/vae/train/video_dc_ae_disc.py并通过--model.from_pretrained指定第一阶段的 checkpoint 继续训练torchrun --nproc_per_node 8 scripts/vae/train.py configs/vae/train/video_dc_ae_disc.py --model.from_pretrained model_ckpt该配置通过_base_ [video_dc_ae.py]继承第一阶段全部配置并在此基础上叠加判别器相关设置discriminator dict( typeN_Layer_discriminator_3D, from_pretrainedNone, input_nc3, n_layers5, conv_clsconv3d ) disc_lr_scheduler dict(warmup_steps0) gen_loss_config dict( gen_start0, disc_weight0.05, ) disc_loss_config dict( disc_start0, disc_loss_typehinge, ) optim_discriminator dict( clsHybridAdam, lr1e-4, eps1e-8, weight_decay0.0, adamw_modeTrue, betas(0.9, 0.98), )其中判别器实现位于 opensora/models/vae/discriminator.pyNLayerDiscriminator3D是一个3D PatchGAN 判别器pix2pix 风格的 3D 扩展默认ndf64、n_layers5、dropout0.30卷积层使用stride2逐步下采样第二层起时间维 stride 为 1、空间维为 2最终输出单通道 logits 图。训练循环见 scripts/vae/train.py中生成器与判别器交替更新先生成器前向 GeneratorLoss反向传播更新生成器再对真实视频与重建视频分别计算判别器 logits用DiscriminatorLoss更新判别器两个优化器与学习率调度器相互独立。3.4 可选wandb 在线监控如果拥有 wandb 账号并希望在线追踪训练进度追加标志--wandb True对应训练脚本中会以cfg.get(wandb_project, Open-Sora)初始化项目基线配置中wandb_project dcae并周期性同步 loss 曲线与梯度范数等指标。四、推理视频重建Inference4.1 权重获取推理前先按 README.md#model-download 中的模型下载指引获取对应权重也可以使用自己训练的模型通过以下标志指定--model.from_pretrained your_model_ckpt_path推理脚本 scripts/vae/inference.py 也支持--ckpt_path覆盖model.from_pretrained。推理过程中会对每个样本执行encode → decode并把原始视频与重建视频分别保存到save_dir/orig与save_dir/recn子目录通过save_sample以fps配置的帧率写出同时累计统计隐特征 z 的逐通道均值与方差并周期性打印。4.2 使用 Video DC-AE 重建torchrun --nproc_per_node 1 --standalone scripts/vae/inference.py configs/vae/inference/video_dc_ae.py --save-dir samples/dcae对应的推理配置 configs/vae/inference/video_dc_ae.py 完整内容如下dtype bf16 batch_size 1 seed 42 dataset dict( typevideo_text, transform_nameresize_crop, fps_max16, data_pathdatasets/pexels_45k_necessary.csv, ) bucket_config { 512px_ar1:1: {96: (1.0, 1)}, } model dict( typedc_ae, model_namedc-ae-f32t4c128, from_pretrained./ckpts/F32T4C128_AE.safetensors, from_scratchTrue, use_spatial_tilingTrue, use_temporal_tilingTrue, spatial_tile_size256, temporal_tile_size32, tile_overlap_factor0.25, ) save_dir samples/video_dc_ae注意推理配置默认按512px、96 帧的 bucket 采样数据进行重建这已经超出模型训练时的 256px/32 帧——正是依靠下文的分块tiling机制才能无退化地处理更大尺寸的输入。4.3 使用 HunyuanVideo VAE 重建Open-Sora 同时把 HunyuanVideo 的 VAE 集成进仓库可用如下命令推理torchrun --nproc_per_node 1 --standalone scripts/vae/inference.py configs/vae/inference/hunyuanvideo_vae.py --save-dir samples/hunyuanvideo_vae对应配置 configs/vae/inference/hunyuanvideo_vae.py 的关键差异点包括model dict( typehunyuan_vae, from_pretrained./ckpts/hunyuan_vae.safetensors, in_channels3, out_channels3, layers_per_block2, latent_channels16, scale_factor0.476986, shift_factor0, use_spatial_tilingTrue, use_temporal_tilingTrue, time_compression_ratio4, )HunyuanVideo VAE 为 16 通道隐空间并带有scale_factor/shift_factor归一化参数其推理 bucket 为 512px、97 帧512px_ar1:1: {97: (1.0, 1)}。该模型实现在 opensora/models/hunyuan_vae 目录下。五、配置文件逐项解读Config Interpretation所有 AE 相关配置统一放在configs/vae/目录下分为训练配置configs/vae/train与推理配置configs/vae/inference两类。训练配置遵循与扩散模型相同的配置规则详见 docs/train.md 中的 Config 一节。5.1 损失配置Loss Config由于 Video DC-AE 基于 DC-AE 架构、没有变分组件训练损失仅由**重建损失reconstruction loss与感知损失perceptual loss**组成。实验表明感知损失权重取0.5效果较好vae_loss_config dict( perceptual_loss_weight0.5, # weigh the perceptual loss by 0.5 kl_loss_weight0, # no KL loss )对照 opensora/models/vae/losses.py 中VAELoss的实现可以更精确地理解这两个参数重建损失为像素级 L1 距离recon_loss l1(video, recon_video)输入与重建帧均按(b t) c h w展平后计算感知损失由冻结的LPIPS网络计算perceptual_loss self.perceptual_loss_fn(video, recon_video)权重由perceptual_loss_weight控制二者合并为 NLL 损失nll_loss recon_loss perceptual_loss * perceptual_loss_weight并经可学习的logvar归一化由于 DC-AE 非变分模型posterior为NoneKL 项恒为 0kl_loss_weight0与之对应。后续阶段加入判别器后AE 的训练损失额外包含生成器generator损失且使用了较小的权重 0.05gen_loss_config dict( gen_start0, # include generator loss from step 0 onwards disc_weight0.05, # weigh the loss by 0.05 )GeneratorLoss见 opensora/models/vae/losses.py的生成器对抗项为g_loss -mean(logits_fake)并通过calculate_adaptive_weight依据生成器最后一层model.get_last_layer()即解码器 project_out 的卷积权重上 NLL 损失与生成器损失的梯度范数比动态计算自适应权重d_weight ||nll_grads|| / (||g_grads|| 1e-4)再乘以disc_weight与gen_start门控因子。gen_start0表示从第 0 步起就启用生成器损失。判别器从零开始训练其损失为简单的hinge 损失disc_loss_config dict( disc_start0, # update the discriminator from step 0 onwards disc_loss_typehinge, # the discriminator loss type )DiscriminatorLoss实现于 opensora/models/vae/losses.py支持三种类型hingemean(ReLU(1 - real)) mean(ReLU(1 fake))的 hinge 形式、vanillasoftplus 形式与wgan-gpWGAN 形式disc_start0意味着判别器从第 0 步就开始更新。5.2 数据桶配置Data Bucket Config官方使用32 帧的 256px 视频训练 Video DC-AEbucket_config { 256px_ar1:1: {32: (1.0, 1)}, }5.3 训练更长帧数或更高分辨率如果训练更长帧数或更高分辨率可以在推理时相应增大spatial_tile_size与temporal_tile_size且不会降低 AE 性能详见下文 Inference Config。这带来的好处是推理例如训练扩散模型时的 AE 前向速度更快代价是 AE 训练变慢。例如把帧数提升到96 帧4 的倍数即可但官方一般推荐使用 32 的倍数bucket_config { 256px_ar1:1: {96: (1.0, 1)}, } grad_checkpoint True或训练更高分辨率如512pxbucket_config { 512px_ar1:1: {32: (1.0, 1)}, } grad_checkpoint True注意此时必须开启梯度检查点grad_checkpoint True以避免 OOM。scripts/vae/train.py中通过set_grad_checkpoint(model)启用DC-AE 编码器/解码器各阶段在forward中经auto_grad_checkpoint包装从而以重计算换取显存。此外判别器训练时若开启了grad_checkpoint需要额外传入--model.disc_off_grad_ckpt True或在配置中直接设置grad_checkpoint True model dict( disc_off_grad_ckpt True, # set to true if your grad_checkpoint is True )这样做的目的是确保自适应损失计算adaptive loss在最后一层有梯度——从 opensora/models/dc_ae/models/dc_ae.py 的 Decoder 实现可以看到disc_off_grad_ckpt会令解码器project_out走普通前向而非梯度检查点包装x self.project_out(x)从而保证get_last_layer()处可获得用于自适应加权的梯度。5.4 推理配置Inference ConfigAE 推理中Open-Sora 把 HunyuanVideo 的tiling分块机制复刻到了 Video DC-AE上可通过以下配置开启model dict( ..., use_spatial_tilingTrue, use_temporal_tilingTrue, spatial_tile_size256, temporal_tile_size32, tile_overlap_factor0.25, ..., )默认情况下空间 tiling 与时间 tiling 均开启以获得最佳性能。由于 Video DC-AE 只在 256px、32 帧的视频上训练过因此spatial_tile_size应设为256、temporal_tile_size应设为32如果你用自己的数据、以其他分辨率或时长训练了 Video DC-AE可以按需调整这两个值。从源码 opensora/models/dc_ae/models/dc_ae.py 可以验证 tiling 的实现细节DCAEConfig中 tiling 相关默认值正是use_spatial_tilingFalse、use_temporal_tilingFalse、spatial_tile_size256、temporal_tile_size32、tile_overlap_factor0.25构造时会校验spatial_tile_size必须能被空间压缩比 32 整除、temporal_tile_size必须能被时间压缩比 4 整除并据此计算隐空间侧的 tile 尺寸spatial_tile_latent_size spatial_tile_size // 32等编码/解码入口encode/decode按时间维度超过temporal_tile_size则走temporal_tiled_encode否则空间维度超过spatial_tile_size则走spatial_tiled_encode的优先级选择分块路径分块时相邻 tile 之间有overlap步长为tile_size * (1 - tile_overlap_factor)重叠区域在隐空间侧通过blend_v/blend_h/blend_t做线性渐变融合避免 tile 边界出现接缝伪影。由此也印证了 5.3 节训练更大分辨率/更长帧数后推理时调大 tile size 不降性能的原理tiling 让模型始终以训练时见过的 256px×32 帧窗口为单位进行编解码再拼接融合成完整大视频。输出样本的保存目录可以用命令行--save_dir your_dir指定也可以在配置中设置例如save_dir ./samples六、结语与后续阅读Video DC-AE 是 Open-Sora 高效视频生成管线的基础设施之一通过 4x时间×32x32空间的高压缩比把扩散模型的 token 数整体减少 16x从底层加速训练与推理配合复刻自 HunyuanVideo 的时空 tiling 机制使得仅在 256px×32 帧上训练的模型也能无损重建更大、更长的视频。本文所涉及的核心代码与配置均可按以下路径在仓库中进一步研读模型实现opensora/models/dc_ae/models/dc_ae.pyDC-AE 编码器/解码器、tiling 与 blend 逻辑损失实现opensora/models/vae/losses.pyVAELoss、GeneratorLoss、DiscriminatorLoss判别器实现opensora/models/vae/discriminator.py3D PatchGAN训练/推理入口scripts/vae/train.py、scripts/vae/inference.py配置文件configs/vae/train 与 configs/vae/inference数据准备与通用训练规则docs/train.md模型下载README.md【免费下载链接】Open-SoraOpen-Sora: Democratizing Efficient Video Production for All项目地址: https://gitcode.com/GitHub_Trending/op/Open-Sora创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
分享:

看完干货,该让你的企业上线了

免费需求沟通 · 48 小时内出具建站方案 · 河南本地可上门