博客

torchcodec：为 PyTorch 提供简单高效的视频解码

作者： 2024年12月11日2025年5月5日暂无评论

我们很高兴正式宣布推出 torchcodec，这是一个用于将视频解码为 PyTorch 张量的库。它速度快、精度高且易于使用。在针对视频运行 PyTorch 模型时，torchcodec 是我们将视频转换为模型可用数据时的推荐方案。

torchcodec 的主要亮点包括：

  • 直观的解码 API，将视频文件视为 Python 的帧序列。我们支持基于索引和基于呈现时间（presentation-time）的帧检索。
  • 对准确性的强调：即使视频具有可变帧率，我们也能确保您获得所请求的正确帧。
  • 丰富的采样 API，使检索帧批次变得既简单又高效。
  • 同类产品中领先的 CPU 解码性能。
  • CUDA 加速解码，支持在同时解码多个视频时实现高吞吐量。
  • 支持您所安装的 FFmpeg 版本中提供的所有编解码器。
  • 针对 Linux 和 Mac 的简单二进制安装方式。

易于使用

简单、直观的 API 是我们的核心设计原则之一。我们可以从简单的解码和提取视频的特定帧开始。

from torchcodec.decoders import VideoDecoder
from torch import Tensor

decoder = VideoDecoder("my_video.mp4")

# Index based frame retrieval.
first_ten_frames: Tensor = decoder[10:]
last_ten_frames: Tensor = decoder[-10:]

# Multi-frame retrieval, index and time based.
frames = decoder.get_frames_at(indices=[10, 0, 15])
frames = decoder.get_frames_played_at(seconds=[0.2, 3, 4.5])

所有解码后的帧均为 PyTorch 张量，可直接输入模型进行训练。

当然，在机器学习训练流水线中，更常见的做法是从视频中采样多个片段。片段只是按呈现顺序排列的帧序列，但这些帧通常不是连续的。我们的采样 API 使这一过程变得简单。

from torchcodec.samplers import clips_at_regular_timestamps

clips = clips_at_regular_timestamps(
  decoder,
  seconds_between_clip_starts=10,
  num_frames_per_clip=5,
  seconds_between_frames=0.2,
)

上述调用会产生一批片段，其中每个片段间隔 10 秒，每个片段包含 5 帧，帧间隔为 0.2 秒。更多详情，请参阅我们关于解码和采样的教程！

极速性能

性能是我们另一个主要的设计原则。机器学习训练的视频解码与用于播放的视频解码有着不同的性能要求。典型的机器学习视频训练流水线会处理许多不同的视频（有时甚至达数百万个！），但每个视频只需采样少量帧（几十到几百帧）。

因此，我们特别关注解码器在视频中多次寻址（seeking）时的性能，即每次寻址后解码少量帧。我们展示了以下四种场景的实验：

  1. 同时解码和转换多个视频中的帧，灵感来自我们在大规模训练流水线数据加载中观察到的场景：a. 十个线程并行解码 50 个视频的批次。
    b. 对每个视频，在均匀间隔的时间点解码 10 帧。
    c. 对每一帧，将其调整为 256×256 分辨率。
  2. 在单个视频中随机位置解码 10 帧。
  3. 在单个视频的均匀间隔时间点解码 10 帧。
  4. 解码单个视频的前 100 帧。

我们将以下视频解码器进行了对比：

  • Torchaudio，仅限 CPU 解码。
  • Torchvision，使用仅限 CPU 解码的 video_reader 后端。
  • Torchcodec，使用 CUDA 进行 GPU 解码。
  • Torchcodec，仅限 CPU 解码。

使用了以下三个视频：

  1. 使用 FFmpeg 的 mandelbrot 生成模式合成的视频。该视频时长 10 秒，60 帧/秒，分辨率 1920×1080。
  2. 与上述相同，但视频时长为 120 秒。
  3. 一段时长 206 秒、29.7 帧/秒、分辨率 960×540 的 NASA 宣传视频。

我们的存储库中提供了实验脚本。我们的实验运行在配备 Intel 处理器（拥有 22 个可用核心）和 NVIDIA GPU 的 Linux 系统上。对于 CPU 解码，所有库均被设置为自动确定最佳线程数。

Benchmark chart

根据实验，我们得出以下结论：

  • 对于我们设计的首要用例（即作为训练数据加载流水线的一部分同时解码多个视频），Torchcodec 的性能表现始终最优。特别是在处理高分辨率视频时，由于解码和转换都在 GPU 上完成，CUDA 带来了显著的性能增益。
  • 在随机采样和均匀采样等频繁寻址的场景下，Torchcodec 在 CPU 上的表现具有竞争力。目前，Torchcodec 在处理较短、文件较小的视频时性能更好。这种性能表现归功于 Torchcodec 对寻址准确性的强调，这涉及初始的线性扫描。
  • 当不需要寻址时（即打开视频文件并从头开始解码），Torchcodec 的竞争力略有下降。这也是由于我们对寻址准确性的强调及初始线性扫描所致。

在 torchcodec 中实现近似寻址模式应能解决这些性能差距，这也是我们在视频解码方面优先级最高的功能。

未来计划

顾名思义，torchcodec 的长远愿景不仅仅是视频解码。我们的下一个重大功能是音频支持——包括从视频中解码音频流，以及解码纯音频媒体。从长远来看，我们希望 torchcodec 成为 PyTorch 的媒体解码库。这意味着随着我们在 torchcodec 中实现相关功能，我们将弃用并最终移除 torchaudio 和 torchvision 中的互补功能。

我们也已规划了视频解码方面的改进，例如前述提到的“近似寻址模式”，供那些愿意以精度换取性能的用户使用。

最重要的是，我们期待社区的反馈！我们非常有兴趣开发社区认为有价值的功能。欢迎来分享您的需求，并影响我们未来的发展方向！