TorchCodec 0.14 发布：CPU 与 CUDA 支持 HDR 视频解码及高速 Wav 解码器

背景

TorchCodec 是 PyTorch 生态中用于高效处理媒体数据（视频、音频）的重要库。随着多模态大模型（Multimodal LLMs）和 AI 视频生成技术的兴起，对媒体数据预处理、解码及格式转换的效率要求日益提高。TorchCodec 旨在通过底层优化，为 PyTorch 用户提供比传统工具（如 FFmpeg 结合 OpenCV）更快速、更原生的媒体解码体验。

此次发布的 TorchCodec 0.14 版本，不仅保持了与 PyTorch 2.11 及更高版本的兼容性，更在音频解码速度和视频动态范围处理上带来了突破性更新，特别是引入了对高动态范围（HDR）视频的原生支持，这标志着该库在专业视频处理领域迈出了关键一步。

核心内容

TorchCodec 0.14 版本主要包含两大核心功能更新：专为 WAV 文件设计的高速解码器，以及对 HDR 视频解码的全面支持。

1. 高速 Wav 解码器 (Fast Wav Decoder)

新版本引入了专用的 WavDecoder，旨在解决音频解码中的性能瓶颈。

• 架构优化：该解码器完全绕过了 FFmpeg 库。传统方案通常依赖 FFmpeg 进行音频解码，而 WavDecoder 直接读取 WAV 文件的数据结构。这种底层直接访问的方式消除了中间层的开销，从而实现了显著的速度提升。
• 格式支持：支持多种采样格式，包括 int16、int32、float32 等，满足不同的精度需求。
• 输入灵活性：支持从本地文件、字节流（bytes）或类文件对象（file-like objects）中读取数据进行解码。

代码示例：

from torchcodec.decoders import WavDecoder

# 初始化解码器
decoder = WavDecoder("audio.wav")

# 获取所有音频样本
# 返回 AudioSamples 对象，包含音频数据和采样率
samples = decoder.get_all_samples()

2. HDR 视频解码支持 (HDR Video Decoding)

VideoDecoder 现在支持高动态范围（HDR）视频的解码，且能够保持数据的完整精度，不再出现色彩信息丢失的情况。

• 精度保留：当指定 output_dtype=torch.float32 时，解码器会输出 RGB 格式的 float32 帧，数值范围在 [0, 1] 之间。这种处理方式保留了完整的 HDR 色彩范围，对于需要高精度色彩处理的 AI 视觉任务至关重要。
• 跨平台支持：该功能同时支持 CPU 和 CUDA 环境，用户可以根据硬件资源灵活选择加速方式。

代码示例：

import torch
from torchcodec.decoders import VideoDecoder

# 初始化 HDR 视频解码器
decoder = VideoDecoder("hdr_video.mp4", output_dtype=torch.float32)

# 获取第一帧
# 返回具有完整 HDR 精度的 float32 数据
frame = decoder[0]

3. 其他改进

• 音频跳转性能提升：AudioDecoder 的 seek（跳转）操作速度大幅优化（参考 Issue #1449），提升了在长音频文件中随机访问特定时间点的效率。
• 移除 NPP 依赖：TorchCodec 不再依赖 NVIDIA 的 NPP (NVIDIA Performance Primitives) 库。这一改动简化了 CUDA 解码环境的安装和使用流程，降低了部署门槛。

关键要点

• 版本兼容性：TorchCodec 0.14 兼容 PyTorch 2.11 及以上版本。
• 音频解码革新：新增 WavDecoder，通过绕过 FFmpeg 直接读取数据，实现比传统方案更快的 WAV 文件解码速度。
• HDR 视频支持：VideoDecoder 新增对 HDR 视频的支持，输出 float32 格式的 RGB 帧（范围 [0, 1]），支持 CPU 和 CUDA，确保色彩精度不丢失。
• 输入多样性：WavDecoder 支持从文件、字节流或类文件对象读取数据。
• 采样格式兼容：音频解码支持 int16、int32、float32 等多种采样格式。
• 性能优化：AudioDecoder 的 seek 操作速度得到显著改善。
• 依赖简化：移除了对 NVIDIA NPP 库的依赖，简化了 CUDA 环境的配置。

意义与影响

TorchCodec 0.14 的发布对 AI 媒体处理工作流具有显著意义：

1. 提升多模态数据处理效率：对于依赖视频和音频输入的多模态大模型，更快的解码速度意味着更短的预处理时间，从而加速训练和推理流程。特别是 WavDecoder 绕过 FFmpeg 的设计，为纯 PyTorch 工作流提供了更轻量、更高效的替代方案。
2. 解锁 HDR 内容潜力：HDR 视频包含更丰富的色彩和亮度信息，对于视频生成、视频理解以及需要高精度色彩分析的应用场景至关重要。原生支持 HDR 解码且保留 float32 精度，使得研究人员和开发者能够直接在 PyTorch 中处理高质量视频素材，无需借助外部工具进行格式转换或精度损失处理。
3. 降低部署复杂度：移除 NPP 依赖是一个重要的工程改进。NPP 库的安装和配置往往较为复杂，且可能引发版本冲突。简化依赖关系使得 TorchCodec 在 CUDA 环境下的部署更加友好，有助于该库在更广泛的 AI 基础设施中被采用。
4. 强化 PyTorch 原生生态：通过提供高性能、原生的媒体解码工具，TorchCodec 进一步巩固了 PyTorch 在多媒体 AI 领域的地位，减少了开发者在不同工具链之间切换的成本。