NVIDIA Cosmos Tokenizer 完全指南:图像与视频神经分词器(Neural Tokenizer)的原理、安装与推理实战
2026/9/15 19:24:54 网站建设 项目流程

NVIDIA Cosmos Tokenizer 完全指南:图像与视频神经分词器(Neural Tokenizer)的原理、安装与推理实战

【免费下载链接】cosmosNVIDIA Cosmos is an open platform of world models, datasets, and tools that enables developers to build Physical AI for robots, autonomous vehicles, smart infrastructure, and more.项目地址: https://gitcode.com/GitHub_Trending/cosmos7/cosmos

导读:本文以 cosmos1/models/tokenizer/README.md 为主体,系统讲解 NVIDIA Cosmos 平台的核心组件Cosmos Tokenizer——一套面向图像与视频的神经分词器套件。它既能把视觉输入压缩为连续潜空间(Continuous Latents),也能量化为离散词元(Discrete Tokens),是构建大规模自回归 Transformer(如 LLM 式视频模型)与扩散生成器的关键基础设施。读完本文,你将掌握其四象限模型选型(CI/DI/CV/DV)、空间 8x/16x 与时间 4x/8x 的压缩机制、JIT 与原生 PyTorch 两种推理模式、CLI 与 Python API 的完整调用方式,并理解 Haar 小波 patching 与 FSQ 有限标量量化等底层实现原理。

一、什么是 Cosmos Tokenizer:从像素到词元的桥梁

Cosmos Tokenizer 是 NVIDIA Cosmos 世界模型平台(面向 Physical AI,服务于机器人、自动驾驶、智慧基础设施等场景)的核心基础组件。它的职责非常纯粹:给定一张图像或一段视频,输出连续的潜在表示(latents)或离散的 token 序列

这套 tokenizer 的定位,正如 README 所述,是推动大规模自回归 Transformer(如 LLM 类模型)或扩散生成器走向可扩展、稳健、高效开发的基石。对开发者而言,它的价值在于:将高维、冗余的原始像素转化为紧凑、规则、便于序列模型建模的表示,从而让后续的世界模型训练与推理成为可能。

Cosmos Tokenizer 共包含4 个模型族,覆盖「图像/视频」×「连续/离散」两个维度:

连续(Continuous, C)离散(Discrete, D)
图像(Images, I)Cosmos-Tokenizer-CICosmos-Tokenizer-DI
视频(Videos, V)Cosmos-Tokenizer-CVCosmos-Tokenizer-DV

其整体架构如下——输入视频先经3D Haar Wavelet处理,进入Causal Encoder(含时间维度的因果卷积与因果注意力),分别产出连续与离散两种潜在空间,再由Causal Decoder解码,最终通过 3D Haar Wavelet 逆变换重建视频:

核心能力:最高 2048x 的总压缩率

README 明确指出 Cosmos Tokenizer 的关键量化指标:

  • 空间压缩率:8x 或 16x;
  • 时间压缩率:4x 或 8x;
  • 总压缩率最高可达 2048x(= 8x × 16x × 16x)。

相比此前的最优方法(SOTA),README 声称 Cosmos Tokenizer 在总压缩率上高出 8 倍,同时保持更优的图像质量,且推理速度快 2x~12x。这些是官方文档声明的项目事实;其具体证据可见后文「评估与性能」章节引用的官方对比图,以及仓库中的推理实现(inference/video_lib.py、inference/image_lib.py)。

二、环境安装:Docker 一键就绪

Cosmos Tokenizer 的推理代码需要 GPU 环境(默认cudabfloat16精度)。官方推荐使用仓库自带的 Docker 镜像,所有命令均在容器内执行。完整步骤见根目录 INSTALL.md:

  1. 安装 NVIDIA Container Toolkit(安装指南仅支持 Ubuntu 24.04 / 22.04 / 20.04);
  2. 克隆仓库并构建、启动容器:
git clone git@github.com:NVIDIA/Cosmos.git cd Cosmos docker build -t cosmos . docker run -d --name cosmos_container --gpus all --ipc=host -it -v $(pwd):/workspace cosmos docker attach cosmos_container

构建完成后,容器内的python3环境已包含torchmediapy等推理所需依赖,可直接进入「下载模型」与「运行推理」环节。

三、下载预训练权重:12 个模型一张表

NVIDIA 在 Hugging Face 上托管了全部12 个 Cosmos-Tokenizer 模型,覆盖 CI/DI/CV/DV 四族在不同压缩率下的组合。模型命名规则为Cosmos-{版本}-Tokenizer-{类型}{压缩率}

模型名称类型压缩率(时间×高×宽)
Cosmos-0.1-Tokenizer-CI8x8连续图像8x8
Cosmos-0.1-Tokenizer-CI16x16连续图像16x16
Cosmos-0.1-Tokenizer-CV4x8x8连续视频4x8x8
Cosmos-0.1-Tokenizer-CV8x8x8连续视频8x8x8
Cosmos-0.1-Tokenizer-CV8x16x16连续视频8x16x16
Cosmos-0.1-Tokenizer-DI8x8离散图像8x8
Cosmos-0.1-Tokenizer-DI16x16离散图像16x16
Cosmos-0.1-Tokenizer-DV4x8x8离散视频4x8x8
Cosmos-0.1-Tokenizer-DV8x8x8离散视频8x8x8
Cosmos-0.1-Tokenizer-DV8x16x16离散视频8x16x16
Cosmos-1.0-Tokenizer-CV8x8x8连续视频(1.0 版)8x8x8
Cosmos-1.0-Tokenizer-DV8x16x16离散视频(1.0 版)8x16x16

使用下面的 Python 片段即可批量下载全部权重(需先在 Hugging Face 生成 Access Token):

from huggingface_hub import login, snapshot_download import os login(token="<YOUR-HF-TOKEN>", add_to_git_credential=True) model_names = [ "Cosmos-0.1-Tokenizer-CI8x8", "Cosmos-0.1-Tokenizer-CI16x16", "Cosmos-0.1-Tokenizer-CV4x8x8", "Cosmos-0.1-Tokenizer-CV8x8x8", "Cosmos-0.1-Tokenizer-CV8x16x16", "Cosmos-0.1-Tokenizer-DI8x8", "Cosmos-0.1-Tokenizer-DI16x16", "Cosmos-0.1-Tokenizer-DV4x8x8", "Cosmos-0.1-Tokenizer-DV8x8x8", "Cosmos-0.1-Tokenizer-DV8x16x16", "Cosmos-1.0-Tokenizer-CV8x8x8", "Cosmos-1.0-Tokenizer-DV8x16x16", ] for model_name in model_names: hf_repo = "nvidia/" + model_name local_dir = "checkpoints/" + model_name print(f"downloading {model_name}...") snapshot_download(repo_id=hf_repo, local_dir=local_dir)

每个 checkpoint 目录checkpoints/{model_name}/下提供三份TorchScript(JIT)模型

├── Cosmos-1.0-Tokenizer-CV8x8x8/ │ ├── encoder.jit │ ├── decoder.jit │ └── autoencoder.jit
  • encoder.jit:仅编码器(像素 → 潜空间 / token);
  • decoder.jit:仅解码器(潜空间 / token → 像素);
  • autoencoder.jit:完整的编码-解码一体模型,可直接做端到端重建。

说明:snapshot_download会下载仓库内全部文件,其中 JIT 模型即上表三件套。下载脚本见仓库根目录 scripts/download_autoregressive.py 的同款模式。

四、命令行推理:图像/视频的编码与重建

仓库为图像与视频分别提供了 CLI 入口,支持glob通配符批量处理。同一套命令对连续/离散分词器都适用,只需把--checkpoint_enc/--checkpoint_dec(或整体--checkpoint)指向对应的 JIT 权重。

4.1 图像自动编码(Autoencoding)

Cosmos-CI8x8 压缩为例,对测试图片(cosmos1/models/tokenizer/test_data/image.png)做「编码 → 解码」重建:

model_name="Cosmos-0.1-Tokenizer-CI8x8" python3 -m cosmos1.models.tokenizer.inference.image_cli \ --image_pattern 'cosmos1/models/tokenizer/test_data/image.png' \ --checkpoint_enc checkpoints/${model_name}/encoder.jit \ --checkpoint_dec checkpoints/${model_name}/decoder.jit

若未指定--output_dir,重建结果默认写到输入目录下的reconstructions子目录,即cosmos1/models/tokenizer/test_data/reconstructions/image.png(该默认行为由 inference/utils.py 中的get_output_filepath实现)。

4.2 视频自动编码

Cosmos-DV4x8x8 压缩为例,对测试视频(cosmos1/models/tokenizer/test_data/video.mp4)做重建:

model_name="Cosmos-0.1-Tokenizer-DV4x8x8" python3 -m cosmos1.models.tokenizer.inference.video_cli \ --video_pattern 'cosmos1/models/tokenizer/test_data/video.mp4' \ --checkpoint_enc checkpoints/${model_name}/encoder.jit \ --checkpoint_dec checkpoints/${model_name}/decoder.jit

默认输出路径为cosmos1/models/tokenizer/test_data/reconstructions/video.mp4

4.3 CLI 参数速查(以 video_cli.py 为准)

参数取值默认值说明
--video_pattern/--image_patternglob 路径输入文件通配符,支持批量
--checkpoint路径None完整 autoencoder JIT 模型
--checkpoint_enc路径None编码器 JIT 模型
--checkpoint_dec路径None解码器 JIT 模型
--modejit/torchjit推理后端,见第五章
--tokenizer_type图像:CI/DI;视频:CV/DVtorch模式需要
--spatial_compression8 或 168空间压缩因子(torch模式用)
--temporal_compression4 或 84时间压缩因子(torch模式用,仅视频)
--short_sizeintNone将短边缩放至该尺寸后再推理
--temporal_windowint17视频按该帧数窗口滑窗处理(仅视频)
--dtypebfloat16bfloat16推理精度
--devicecudacuda运行设备
--output_dir路径None输出目录,默认<输入目录>/reconstructions
--output_fpsfloat24.0输出视频帧率(仅视频)
--save_inputflag关闭同时保存输入副本(_input后缀)

注意:当--checkpoint_enc--checkpoint_dec--checkpoint三者都未提供时,CLI 会中止并提示需要编码器/解码器或完整 autoencoder(见 image_cli.py)。

五、两种推理后端:JIT 与原生 PyTorch

默认情况下 CLI 加载encoder.jit/decoder.jit这类TorchScript 编译模型直接推理,无需任何网络定义代码。若要运行原生 PyTorch 模式,在命令后追加--mode=torch即可:

  • PyTorch 模式下,模型会根据仓库内的原生网络定义脚本(networks/)重新实例化,因此必须额外提供参数来指定网络结构;
  • 权重仍来自 JIT 文件--checkpoint_enc--checkpoint_dec--checkpoint依旧指向 JIT 路径,加载时会从 JIT 模型中提取state_dict来初始化原生 PyTorch 模型。

5.1 图像:原生 PyTorch 实例化Cosmos-DI(空间压缩 8x)

需要追加的三个参数:

  • --mode=torch
  • --tokenizer_type=DI
  • --spatial_compression=8
model_name="Cosmos-0.1-Tokenizer-DI8x8" python3 -m cosmos1.models.tokenizer.inference.image_cli \ --image_pattern 'cosmos1/models/tokenizer/test_data/*.png' \ --mode=torch \ --tokenizer_type=DI \ --spatial_compression=8 \ --checkpoint_enc checkpoints/${model_name}/encoder.jit \ --checkpoint_dec checkpoints/${model_name}/decoder.jit

5.2 视频:原生 PyTorch 实例化Cosmos-CV(8x8x8)

需要追加四个参数:

  • --mode=torch
  • --tokenizer_type=CV
  • --temporal_compression=8
  • --spatial_compression=8
model_name="Cosmos-1.0-Tokenizer-CV8x8x8" python3 -m cosmos1.models.tokenizer.inference.video_cli \ --video_pattern 'cosmos1/models/tokenizer/test_data/*.mp4' \ --mode=torch \ --tokenizer_type=CV \ --temporal_compression=8 \ --spatial_compression=8 \ --checkpoint_enc checkpoints/${model_name}/encoder.jit \ --checkpoint_dec checkpoints/${model_name}/decoder.jit

5.3 PyTorch 模式背后的实现

从源码看,torch模式的实质是:CLI 根据tokenizer_type从 networks/init.py 的TokenizerConfigs枚举中取出默认配置字典,再用命令行参数覆盖spatial_compression/temporal_compression;随后 inference/utils.py 中的_load_pytorch_model依据name字段从TokenizerModels枚举(CI→ContinuousImageTokenizer、DI→DiscreteImageTokenizer、CV→CausalContinuousVideoTokenizer、DV→CausalDiscreteVideoTokenizer)实例化网络,最后把 JIT 权重load_state_dict(strict=False)灌入。这意味着你完全可以用这套配置机制构造自己的自定义压缩率网络。

六、Python API 实战:编码为连续潜空间 / 离散 token

除 CLI 外,仓库提供面向库调用的 Python 类:图像用ImageTokenizer(inference/image_lib.py),视频用CausalVideoTokenizer(inference/video_lib.py)。两者均以torch.nn.Module形式封装,支持传入完整 autoencoder 或独立的 encoder/decoder JIT 路径,默认运行在cudabfloat16

6.1 编码为连续潜空间(Continuous Latents)

import torch from cosmos1.models.tokenizer.inference.video_lib import CausalVideoTokenizer model_name = "Cosmos-0.1-Tokenizer-CV4x8x8" input_tensor = torch.randn(1, 3, 9, 512, 512).to('cuda').to(torch.bfloat16) # [B, C, T, H, W] encoder = CausalVideoTokenizer(checkpoint_enc=f'checkpoints/{model_name}/encoder.jit') (latent,) = encoder.encode(input_tensor) torch.testing.assert_close(latent.shape, (1, 16, 3, 64, 64)) # 解码重建 decoder = CausalVideoTokenizer(checkpoint_dec=f'checkpoints/{model_name}/decoder.jit') reconstructed_tensor = decoder.decode(latent) torch.testing.assert_close(reconstructed_tensor.shape, input_tensor.shape)

这里的latent形状为(1, 16, 3, 64, 64),语义拆解如下:

  • 输入[B=1, C=3, T=9, H=512, W=512],9 帧 512x512 RGB 视频;
  • 时间维度 9 → 3:时间压缩 4x 对应1 + (T-1)/4 = 3个时间潜位置(第一个潜位置代表首帧,即因果式对齐,源码注释见 video_lib.py);
  • 空间维度 512 → 64:8x8 空间压缩;
  • C=16:连续潜空间的通道数(该常量在 networks/configs.py 的latent_channels=16/z_channels=16中定义)。

6.2 编码为离散 token(Discrete Tokens)

import torch from cosmos1.models.tokenizer.inference.video_lib import CausalVideoTokenizer model_name = "Cosmos-0.1-Tokenizer-DV4x8x8" input_tensor = torch.randn(1, 3, 9, 512, 512).to('cuda').to(torch.bfloat16) # [B, C, T, H, W] encoder = CausalVideoTokenizer(checkpoint_enc=f'checkpoints/{model_name}/encoder.jit') (indices, codes) = encoder.encode(input_tensor) torch.testing.assert_close(indices.shape, (1, 3, 64, 64)) torch.testing.assert_close(codes.shape, (1, 6, 3, 64, 64)) # 解码重建 decoder = CausalVideoTokenizer(checkpoint_dec=f'checkpoints/{model_name}/decoder.jit') reconstructed_tensor = decoder.decode(indices) torch.testing.assert_close(reconstructed_tensor.shape, input_tensor.shape)

离散模式的输出有两个:

  • indices:形状(1, 3, 64, 64),元素为[1..64K]区间的整数值——64K(65536)即 FSQ 隐式码本大小(由 levels(8,8,8,5,5,5)相乘得到:8×8×8×5×5×5 = 64000,README 以 64K 约称)。第一个整数图对应首帧;
  • codes:量化前的连续潜变量,形状(1, 6, 3, 64, 64),其中C=6FSQ 的 level 维度数(6 个标量量化通道)。

6.3 端到端重建与任意长度视频

  • 若构造时传入checkpoint(完整 autoencoder JIT),autoencode()会直接跑完整模型;否则内部自动串联encode → decode(见 video_lib.py)。
  • forward(video, temporal_window=17)支持任意长度视频的滑动窗口重建:按temporal_window逐窗处理,先空间补零 + 时间边缘反射 padding 到对齐尺寸,推理后再裁剪回原尺寸并拼接(实现见 inference/utils.py 的pad_video_batch/unpad_video_batch)。

七、源码级原理:压缩率从哪来

7.1 Haar 小波 patching:无参数下采样

图像/视频在进入主干网络前,先经过Patcher做可逆下采样(modules/patching.py):

  • 2D 图像使用 Haar 小波(_WAVELETS = {"haar": [0.7071, 0.7071], ...}),每做一次离散小波变换(DWT),把通道从c扩到4c、空间尺寸减半——例如[3, 256, 256] → [12, 128, 128] → [48, 64, 64](源码 docstring 中的示例);
  • 3D 视频使用Patcher3D,每次 DWT 将通道扩到 8 倍、三个维度各减半,同时对首帧做 repeat_interleave以维持因果对齐;
  • 对应的UnPatcher/UnPatcher3D用转置卷积完成逆小波变换恢复像素。patch_size=4时叠加两次 DWT,正好贡献 4x 的空间下采样。

这正是「patch_method='haar'」配置(见 networks/configs.py 中各 dict 的patch_size=4, patch_method="haar")的含义——小波变换完全无学习参数,比可学习的 stride 卷积更高效且可逆。

7.2 FSQ 有限标量量化:无需码本查找的离散化

离散 tokenizer(DI/DV)使用FSQ(Finite Scalar Quantization)而非传统 VQ 码本查找(modules/quantizers.py):

  • 每个潜在向量的 6 个通道分别按levels=[8,8,8,5,5,5]的整数级数独立量化(bound中用 tanh 约束范围,round_ste保留直通梯度);
  • 隐式码本大小 = 各级数之积 = 64000(约 64K),索引由codes_to_indices用混合进制_basis计算得到,解码时由indices_to_codes反向还原;
  • 由于无需查询嵌入表,FSQ 天然支持超大批次与高吞吐,这也是官方宣称推理更快的原因之一。

7.3 因果视频 tokenizer:首帧先行的时序设计

视频族(CV/DV)的关键在于Causal(因果)设计:

  • 时间维度压缩通过因果卷积与因果注意力实现,保证解码时任一帧只依赖其之前(含自身)的帧,从而支持流式生成;
  • 时间压缩 T→t 满足t = 1 + (T-1)/压缩率,例如 9 帧 + 4x 时间压缩 → 3 个时间潜位置,首帧潜位置与首帧严格对齐(README 与 video_lib.py 的 docstring 均明确此约定)。

7.4 默认网络配置速查

networks/configs.py 中四个族的关键超参如下:

配置项CIDICVDV
channels(基础通道数)128128128128
channels_mult(各分辨率通道倍增)[2,4,4][2,4,4][2,4,4][2,4,4]
num_res_blocks2222
patch_size/patch_method4 / haar4 / haar4 / haar4 / haar
默认空间压缩1616816
默认时间压缩88
z_channels/latent_channels16 / 16256 / —16 / 1616 / —
quantizerFSQFSQ
embedding_dim/levels6 / [8,8,8,5,5,5]6 / [8,8,8,5,5,5]

提示:表中为各族的默认压缩率,CLI 的--spatial_compression/--temporal_compression可在 PyTorch 模式下覆盖默认值,但必须与所下载权重的实际压缩率一致,否则加载的state_dict无法匹配。

八、NeMo 集成:JIT 推理与数据集 tokenization

除了仓库自带的推理代码,Cosmos Tokenizer 还深度集成到 NVIDIA NeMo 框架(nemo/README.md),支持:

  • JIT 推理:安装 NeMo(GitHub main 分支)后,直接用CausalVideoTokenizer.from_pretrained(model_name)一行加载模型:
import torch from nemo.collections.common.video_tokenizers.cosmos_vision_tokenizer import CausalVideoTokenizer model_name = "Cosmos-0.1-Tokenizer-CV4x8x8" model = CausalVideoTokenizer.from_pretrained(model_name) input_tensor = torch.randn(1, 3, 9, 512, 512).to('cuda').to(torch.bfloat16) (latent, ) = model.encode(input_tensor)
  • 数据集 tokenization 与多模态训练:NeMo 提供了用 Cosmos Tokenizer 构造多模态训练数据集的示例,可直接对接后续的世界模型训练流程;
  • NeMo 后训练(Post-training):NeMo 框架目前支持对Cosmos-1.0-Tokenizer-CV8x8x8Cosmos-1.0-Tokenizer-DV8x16x16两个模型做领域微调(尤其针对自动驾驶数据),推荐在 H100-80GB 或 A100-80GB 上运行,具体步骤见 nemo/README.md。

九、评估与性能:官方基准结果

9.1 DAVIS 数据集量化对比

README 报告了在 DAVIS(Perazzi et al., 2016)数据集上与历史 tokenizer 的定量对比:即使在更高的压缩率(8x8x8 与 8x16x16)下,Cosmos Tokenizer 仍优于此前方法,展现了出色的「压缩-质量」权衡:

9.2 参数量与推理耗时

README 还对比了各模型在单张 A100 80GB GPU 上的参数量与每图/每帧平均编码、解码耗时:Cosmos Tokenizer 相比此前方法快 2x~12x,且模型尺寸更小,体现了极高的 tokenization 效率:

以上数据均为 README 官方声明,复现与进一步验证可参考官方公开的 TokenBench 基准(见下节)及仓库内的评估代码。

十、TokenBench:视频 tokenizer 的标准化评测基准

为标准化 Cosmos Tokenizer 的评估,官方还配套发布了TokenBench——一个覆盖机器人操作、驾驶、第一人称视角(egocentric)、网络视频等多领域的综合基准,包含高分辨率、长时长视频,专为评测视频 tokenizer 设计。TokenBench 面向社区公开,可用于与其他视频 tokenizer 在同一标准下横向比较压缩质量与效率。

十一、许可证与引用

  • 模型权重:采用 NVIDIA Open Model License,该协议确认:模型可商用;允许创建和分发衍生模型;NVIDIA 不对使用模型(或衍生模型)生成的任何输出主张所有权。
  • 仓库代码:本仓库(含 tokenizer 代码)采用Apache 2.0许可,见根目录 LICENSE。

若在你的工作中使用了 Cosmos Tokenizer,请按如下方式引用:

@article{agarwal2025cosmos, title={Cosmos World Foundation Model Platform for Physical AI}, author={NVIDIA et. al.}, journal={arXiv preprint arXiv:2501.03575}, year={2025} }

十二、致谢与源码归属

仓库注明,modules/ 目录中的部分代码实现派生自以下开源项目:CompVis/stable-diffusion、lucidrains/magvit2-pytorch、lucidrains/vector-quantize-pytorch、CompVis/taming-transformers。这也解释了为什么 modules/quantizers.py 中 FSQ/LFQ/VQ 的实现与上述项目保持接口与算法的一致性——理解这些上游约定,有助于你在二次开发时快速定位量化器行为。

总结

Cosmos Tokenizer 以「连续/离散 × 图像/视频」四象限覆盖视觉 tokenization 的全部主流需求,通过 Haar 小波 patching、因果式 3D 网络与 FSQ 量化,把最高 2048x 的总压缩率、64K 隐式码本与高效推理组合在一起。无论是想快速用 CLI 做图像/视频重建、用 Python API 提取潜空间或离散 token,还是打算对接 NeMo 做数据集构建与领域后训练,上文给出的模型选型表、参数速查表与源码级原理分析都已覆盖;下一步即可按 INSTALL.md 搭建环境、下载对应 checkpoint 并跑通第一条推理命令。

【免费下载链接】cosmosNVIDIA Cosmos is an open platform of world models, datasets, and tools that enables developers to build Physical AI for robots, autonomous vehicles, smart infrastructure, and more.项目地址: https://gitcode.com/GitHub_Trending/cosmos7/cosmos

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询