作者：PyTorch 团队

我们很高兴正式宣布 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 训练管线中更常见的是从视频中采样多个片段（clips）。片段只是按演示顺序排列的一系列帧——但帧通常不是连续的。我们的采样 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 视频训练管线会处理许多不同的视频（有时甚至数百万个！），但每个视频只采样少量帧（几十到几百帧）。

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

  1. 同时解码和转换来自多个视频的帧，灵感来自于我们在大规模训练管线的数据加载中看到的需求

    a. 十个线程并行解码每批 50 个视频。
    b. 对于每个视频，在均匀间隔的时间点解码 10 帧。
    c. 对于每一帧，将其大小调整为 256x256 分辨率。

  2. 在单个视频的随机位置解码 10 帧。
  3. 在单个视频的均匀间隔时间点解码 10 帧。
  4. 解码单个视频的前 100 帧。

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

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

使用以下三个视频

  1. 使用 FFmpeg 的 mandelbrot 生成模式合成生成的视频。该视频长 10 秒，帧率为 60 帧/秒，分辨率为 1920x1080。
  2. 与上同，但视频长 120 秒。
  3. 一段来自 NASA 的宣传视频，长 206 秒，帧率为 29.7 帧/秒，分辨率为 960x540。

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

Benchmark chart

从我们的实验中，我们得出几个结论

  • 对于我们设计它的主要用例——作为训练数据加载管线的一部分同时解码许多视频——torchcodec 始终是性能最佳的库。特别是高分辨率视频，在使用 CUDA 时性能提升显著，因为解码和转换都在 GPU 上进行。
  • 在 CPU 上，对于大量寻址的用例（例如随机采样和均匀采样），torchcodec 具有竞争力。目前，torchcodec 在文件较短、文件大小较小的视频上性能更好。这种性能差异是由于 torchcodec 强调寻址准确性，这涉及一次初始线性扫描。
  • 在没有寻址的情况下，即打开视频文件并从头开始解码，torchcodec 则不具有竞争力。这同样是由于我们对寻址准确性和初始线性扫描的强调。

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

下一步是什么？

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

我们还有一系列视频解码改进计划，例如之前提到的近似寻址模式，适用于愿意牺牲准确性以换取性能的用户。

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