由 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 张量，可以直接用于模型训练。

当然，在机器学习训练流程中，从视频中采样多个片段更为常见。片段只是按展示顺序排列的帧序列——但这些帧通常不是连续的。我们的采样 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 秒。请参阅我们关于解码和采样的教程以获取更多信息！

快速性能

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

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

  1. 受到我们在大规模训练管道数据加载中看到的启发，我们同时解码和转换多个视频的帧：

    a. 使用十个线程并行解码 50 个视频的批次。
    b. 对于每个视频，以均匀的时间间隔解码 10 帧。
    对于每一帧，将其调整至 256x256 分辨率。

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

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

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

使用以下三个视频：

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

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

Benchmark chart

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

  • Torchcodec 是我们为其设计的首要用途中表现最稳定的库：作为训练数据加载管道的一部分，同时解码多个视频。特别是，在 CUDA 下，解码和转换都在 GPU 上执行时，高分辨率视频的解码效果和转换效果都有显著提升。
  • Torchcodec 在 CPU 上对随机和均匀采样等寻求密集型用例具有竞争力。目前，torchcodec 的性能在较短的视频和较小的文件大小上表现更佳。这种性能得益于 torchcodec 对寻求准确性的重视，这涉及到初始的线性扫描。
  • 当没有寻求操作时，即从开头打开视频文件进行解码，torchcodec 的竞争力并不强。这同样是因为我们重视寻求准确性以及初始的线性扫描。

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

接下来是什么？

如其名所示，torchcodec 的长期未来不仅仅是视频解码。我们的下一个重大特性是音频支持——从视频和仅音频媒体中解码音频流。从长远来看，我们希望 torchcodec 成为 PyTorch 的媒体解码库。这意味着随着我们在 torchcodec 中实现功能，我们将弃用并最终从 torchaudio 和 torchvision 中移除互补特性。

我们还准备了视频解码的改进，例如之前提到的近似搜索模式，对于那些愿意牺牲准确性以换取性能的用户。

最重要的是，我们希望得到社区反馈！我们最感兴趣的是开发社区认为有价值的特性。来分享您的需求，影响我们的未来方向！