magvit2-pytorch API速查手册VideoTokenizer与Trainer核心接口完全指南【免费下载链接】magvit2-pytorchImplementation of MagViT2 Tokenizer in Pytorch项目地址: https://gitcode.com/gh_mirrors/ma/magvit2-pytorch本手册为新手与普通用户整理了 magvit2-pytorchMagViT2 视频分词器的 Pytorch 实现最核心的VideoTokenizer与VideoTokenizerTrainerAPI 速查指南。MagViT2 是当前视频生成/理解领域 SOTA 的 Tokenizer 方案其核心思想是把视频压缩成离散 token 序列为后续语言模型式生成铺路。本文用最少的代码、最快的路径带你掌握两大核心类的全部常用接口。上图展示了不同 tokenizer 的图像重建效果对比分辨率、LPIPS 越低越好直观体现 MagViT2 分词器的重建能力。一、快速安装与最小示例pip install magvit2-pytorch训练一个基础 tokenizer 只需两步创建VideoTokenizer模型再交给VideoTokenizerTrainer。完整示例可参考仓库根目录的 README.md对应源码入口见 magvit2_pytorch/init.py。from magvit2_pytorch import VideoTokenizer, VideoTokenizerTrainer tokenizer VideoTokenizer( image_size 128, # 输入分辨率 init_dim 64, # 起始通道数 max_dim 512, # 最大通道数 codebook_size 1024, # 码本大小 layers ( # 网络层配置详见下文 residual, compress_space, (consecutive_residual, 2), compress_time, attend_time, ) ) trainer VideoTokenizerTrainer( tokenizer, dataset_folder /path/to/dataset, # 视频或图片文件夹 dataset_type videos, # videos 或 images batch_size 4, learning_rate 2e-5, num_train_steps 100000 ) trainer.train()二、VideoTokenizer 核心参数速查VideoTokenizer定义于 magvit2_pytorch/magvit2_pytorch.py以下参数最常用参数默认值作用说明image_size必填输入视频帧分辨率宽高相同layers(residual,)*3编码器/解码器层结构支持residual、compress_space、compress_time、attend_space、linear_attend_space、attend_time等codebook_sizeNoneLFQ 码本大小use_fsqFalse时必填use_fsqFalse是否改用 FSQ 有限标量量化器fsq_levelsNoneFSQ 每维级别列表启用时替代codebook_sizenum_codebooks1码本数量多个码本各带独立熵正则channels3输入通道数RGB 视频为 3init_dim64首层卷积通道数max_diminf压缩层最大通道数上限attn_heads/attn_dim_head8/32空间/时间注意力头数与头维度flash_attnTrue是否使用 Flash Attention 加速use_ganTrue是否启用判别器对抗训练perceptual_loss_weight1e-1VGG 感知损失权重adversarial_loss_weight1.对抗损失权重separate_first_frame_encodingFalse是否单独编码视频首帧 快速记忆layers控制网络骨架codebook_size控制压缩率*_loss_weight控制训练目标占比。三、VideoTokenizer 核心方法速查所有方法均定义于 magvit2_pytorch/magvit2_pytorch.py1. 前向推理 forward —— 一个函数四种模式通过开关参数控制返回值且return_loss、return_codes、return_discr_loss三者至多同时启用一个recon tokenizer(video) # 默认返回重建视频 codes tokenizer(video, return_codes True) # 返回离散 token 码 codes, recon tokenizer(video, return_codes True, return_recon True) loss, breakdown tokenizer(video, return_loss True) # 生成器总损失 discr_loss, d_break tokenizer(video, return_discr_loss True) # 判别器损失2. 编码与解码三件套encoded tokenizer.encode(video) # 编码特征可量化 quantized tokenizer.decode(quantized) # 解码回视频 video_recon tokenizer.decode_from_code_indices(codes) # 由 token 索引解码 codes tokenizer.tokenize(video) # 最常用一键离散化其中decode_from_code_indices会自动处理(b, f*h*w)扁平化索引与(b, f, h, w)特征图的相互转换无需手动 reshape。3. 模型存取三连tokenizer.save(./tokenizer.pt) # 保存权重配置 tokenizer.load(./tokenizer.pt) # 从 checkpoint 加载 model VideoTokenizer.init_and_load_from(./tokenizer.pt) # 一步初始化加载checkpoint 中会自动保存构造配置pickle序列化于_configs因此init_and_load_from无需手动传入参数即可还原模型。四、VideoTokenizerTrainer 核心接口速查VideoTokenizerTrainer定义于 magvit2_pytorch/trainer.py负责数据加载、优化器、EMA、分布式训练与断点续训。1. 必填参数与最常用选项参数默认值作用说明model必填传入VideoTokenizer实例batch_size必填批大小num_train_steps必填总训练步数dataset_folderNone数据文件夹路径dataset_typevideosvideos或images论文表明先用图片预训练对视频合成有效num_frames17视频数据集抽帧数learning_rate1e-5学习率grad_accum_every1梯度累积步数checkpoints_folder./checkpoints断点保存目录results_folder./results验证 GIF 输出目录validate_every_step100每 N 步验证一次checkpoint_every_step100每 N 步保存一次断点use_wandb_trackingFalse是否启用 wandb 实验追踪2. 训练与追踪三件套trainer.train() # 一键启动训练含验证与断点 trainer.save(./checkpoint.pt) # 保存完整训练状态 trainer.load(./checkpoint.pt) # 恢复训练含优化器/调度器/步数wandb 追踪需要先开启use_wandb_trackingTrue再配合trackers上下文管理器使用with trainer.trackers(project_name magvit2, run_name baseline): trainer.train()3. 训练后必用EMA 模型训练器内部维护了指数移动平均EMA模型推理时务必使用 EMA 版本效果通常更稳定ema_tokenizer trainer.ema_tokenizer # 获取 EMA 模型 codes ema_tokenizer.tokenize(video) # (1, 9, 16, 16) 时间4x/空间8x下采样 recon ema_tokenizer.decode_from_code_indices(codes) trainer.tokenize(video) # 快捷方式等价于上两行五、常见问题与避坑指南帧数必须对齐下采样因子输入帧数减 1 后必须能被时间下采样倍数整除否则前向会断言报错例如示例配置时间下采样 4 倍则17 1 4*4帧刚好匹配。LFQ 与 FSQ 二选一use_fsqFalse时必须设置codebook_sizeuse_fsqTrue时必须设置fsq_levels且有效码本大小等于各 level 的累乘。空间压缩受image_size限制每层compress_space会把特征图减半因此layers中空间压缩层数量受输入分辨率约束。对抗训练有预热可通过discr_start_after_step控制判别器在训练若干步后再启用避免早期训练不稳定。设备管理模型会自动注册device属性无需手动.to(device)操心多卡由 accelerate 自动处理。六、结语至此magvit2-pytorch 的核心 API 已全部覆盖VideoTokenizer 负责视频↔token的编解码与损失计算VideoTokenizerTrainer 负责训练闭环与 EMA 推理。搭配本文的速查表你可以轻松完成从数据准备、模型配置、训练到离散化推理的全流程。若需深入源码编码器/解码器与量化器实现位于 magvit2_pytorch/magvit2_pytorch.py训练与数据管线分别位于 magvit2_pytorch/trainer.py 与 magvit2_pytorch/data.py祝你在视频生成的道路上一帆风顺【免费下载链接】magvit2-pytorchImplementation of MagViT2 Tokenizer in Pytorch项目地址: https://gitcode.com/gh_mirrors/ma/magvit2-pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考