跳转到主要内容
博客

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

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

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

torchcodec 的亮点包括：

  • 直观的解码 API，将视频文件视为帧的 Python 序列。我们支持基于索引和基于演示时间的帧检索。
  • 强调准确性：我们确保您获得请求的帧，即使您的视频具有可变帧速率。
  • 丰富的采样 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 张量，可以直接输入模型进行训练。

当然，在 ML 训练管道中更常见的是从视频中采样多个剪辑。剪辑只是按演示顺序排列的一系列帧，但这些帧通常并*不*连续。我们的采样 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 秒。有关更多信息，请参阅我们关于解码和采样的教程！

快速性能

性能是我们的另一个主要设计原则。为 ML 训练解码视频与为播放解码视频有不同的性能要求。典型的 ML 视频训练管道将处理许多不同的视频（有时达到数百万！），但只从每个视频中采样少量帧（几十到几百）。

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

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

我们比较了以下视频解码器：

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

使用以下三个视频：

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

该实验脚本位于我们的代码库中。我们的实验在配备 Intel 处理器（22 个可用核心）和 NVIDIA GPU 的 Linux 系统上运行。对于 CPU 解码，所有库都被指示自动确定要使用的最佳线程数。

Benchmark chart

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

  • 对于我们设计它的主要用例，torchcodec 始终是性能最佳的库：作为训练数据加载管道的一部分，同时解码许多视频。特别是，高分辨率视频在使用 CUDA 时会获得巨大的提升，因为解码和转换都发生在 GPU 上。
  • 在 CPU 上，对于随机采样和均匀采样等需要大量查找的用例，torchcodec 具有竞争力。目前，torchcodec 在文件大小较短的视频上表现更好。这种性能得益于 torchcodec 对查找准确性的重视，这涉及到初始的线性扫描。
  • 当没有查找时，即打开视频文件并从头开始解码时，torchcodec 的竞争力不如其他库。这同样是由于我们对查找准确性和初始线性扫描的重视。

在 torchcodec 中实现近似查找模式应该可以解决这些性能差距，这是我们视频解码的最高优先级功能。

下一步是什么？

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

我们还有视频解码的改进，例如前面提到的近似查找模式，适用于那些愿意牺牲准确性以换取性能的用户。

最重要的是，我们正在寻求社区的反馈！我们最感兴趣的是开发社区认为有价值的功能。请分享您的需求并影响我们未来的方向！