PyTorch 媒体处理栈演进：TorchCodec 统一编解码，TorchVision 与 TorchAudio 专注变换

TL;DR

如果你需要解码或编码媒体，无论是图像、视频还是音频，请使用 TorchCodec。如果你需要转换媒体，图像和视频请使用 TorchVision，音频请使用 TorchAudio。这些库中曾经包含的其他所有内容（模型、数据集等），如今由更广泛的生态系统更好地提供服务，例如 HuggingFace 系列库。

图像、视频和音频如今已成为许多模型开发的核心，从视觉语言模型到生成图像和视频的扩散模型皆是如此。训练这些模型意味着将媒体解码为张量并对其进行转换，而生成模型随后需要将其输出编码回媒体文件。在过去两年中，我们将 PyTorch 的媒体技术栈整合到三个库中：TorchCodec、TorchVision 和 TorchAudio。本文介绍了最终的整合结果及其原因。

过去与现在

TorchCodec：解码与编码的唯一去处

几年前，媒体解码和编码能力分散在 TorchVision 和 TorchAudio 中，且存在部分重复，各自构建了自己的技术栈。TorchVision 有多个入口点（io.read_video() 和 io.VideoReader()），分布在三种后端上：Python 中的 PyAV、C++ FFmpeg 后端以及 CUDA/NVCUVID 后端。只有 PyAV 后端开箱即用，因为另外两个后端需要从源码构建，并且与特定的 FFmpeg 版本绑定。TorchAudio 有用于视频和音频的 StreamReader 和 StreamWriter，以及带有各种后端的额外音频解码工具（FFmpeg、libsoundfile、libsox）。图像解码则位于 TorchVision 中。

针对这一情况，我们将媒体的所有解码和编码能力整合到 TorchCodec 中：图像、视频和音频，在 CPU 和 CUDA 上均可使用。过去位于 TorchVision 或 TorchAudio 中的 API 已被弃用或移除。我们这样做有三个原因：

为用户提供统一的集中位置。所有媒体 I/O 都集中在 TorchCodec 中，用户不再需要弄清楚哪个库恰好提供了他们所需的解码器，而且他们得到的是一个整体设计的库，而不是一组相互独立的 API。

性能优化。单一的集中位置也意味着我们的开发工作集中在一个库上，而不是分散在多个库中。当解码和编码存在于多个实现中时，很难判断哪一个值得进行优化工作。现在这一点很明确，我们在 TorchCodec 中加速的任何内容都会让所有人受益。TorchCodec 通常比我们之前在 TorchVision 和 TorchAudio 中的实现性能更好，尤其是在 CUDA 视频解码方面。

维护与依赖。媒体 I/O 涉及大量依赖：六个主要的 FFmpeg 版本（撰写本文时为 4 到 9）、用于 GPU 解码的 NVIDIA 编解码器 SDK，以及每种图像格式对应一个库（libjpeg、libpng 等），每个库都有各自关于我们可以发布哪些内容的许可规则。除此之外，它们都不是纯 Python 依赖：它们全都是 C/C++ 库，这使得构建、CI 和发布变得更加困难。所有这些复杂性现在都集中在 TorchCodec 中，这大大减轻了 TorchVision 和 TorchAudio 的维护负担。

TorchVision 和 TorchAudio：现在专注于转换媒体

TorchVision 和 TorchAudio 曾支持广泛的产品面，涵盖模型、数据集、流水线、变换以及上述 I/O。如今，这两个库都聚焦于它们最擅长的部分：变换。其余部分不再处于积极开发中。

这里的驱动因素是观察到 TorchVision 和 TorchAudio 中的变换是社区中最活跃的使用领域。相比之下，模型、数据集和流水线的使用量滞后，因为出现了优秀的替代方案，例如 HuggingFace 库。缩小 TorchAudio 和 TorchVision 的范围使我们能够加倍投入它们的优势，以及生态系统中其他部分未覆盖的部分。

这一转变对 TorchAudio 尤其具有破坏性：许多 API 被弃用并最终移除。我们感谢社区的耐心和反馈，这使我们能够重新考虑最初的计划。我们原本计划移除的几个流行 API 被保留了下来。这次迁移很困难，但我们认为这是值得的。更小的 TorchAudio 才是真正能够维护的，因此用户也可以继续依赖它。

它们如何协同工作

在实践中，分工很简单：TorchCodec 将媒体文件或编码字节转换为张量（解码），并将张量转换回媒体文件（编码）。TorchVision 和 TorchAudio 则对中间的张量进行变换。

对于视频，使用 TorchCodec 的 VideoDecoder 和 TorchVision 的 v2 transforms：

import torch
from torchcodec.decoders import VideoDecoder
from torchvision.transforms import v2

clip = VideoDecoder("video.mp4")[10:20]   # FrameBatch, uint8, (10, C, H, W)

transform = v2.Compose([
    v2.RandomResizedCrop(224),
    v2.ToDtype(torch.float32, scale=True),
])
batch = transform(clip.data)

对于音频，使用 TorchCodec 的 AudioDecoder 和 TorchAudio 的 transforms：

from torchcodec.decoders import AudioDecoder
from torchaudio.transforms import MelSpectrogram

samples = AudioDecoder("audio.mp3").get_all_samples()   # AudioSamples

mel = MelSpectrogram(sample_rate=samples.sample_rate)(samples.data)

图像的工作方式相同：decode_image 及其同类返回纯张量，直接进入 torchvision.transforms.v2。编码则朝相反方向进行，使用 VideoEncoder、AudioEncoder、JpegEncoder 和 PngEncoder。

在发布方面，这三个库现在都是 ABI 稳定的。这意味着给定版本的 TorchCodec、TorchVision 或 TorchAudio 不再绑定到单一版本的 PyTorch：它可以继续与之后的 PyTorch 版本一起工作。它们不再需要为每个 PyTorch 版本重新构建，因此它们的发布节奏不再与 PyTorch 匹配。这是预期之中的！

如果你仍在使用 TorchVision 或 TorchAudio 中的解码或编码 API，现在是迁移到 TorchCodec 的好时机：从我们的迁移指南开始，如果你需要的内容缺失，请在 GitHub 上告诉我们。

关于媒体处理库还有很多要说的，我们将在后续文章中介绍：每个库的新内容，以及如何从完整的解码和变换流水线中获得最佳性能。敬请期待！

文章 Evolution of the PyTorch Media Processing Landscape 首次出现在 PyTorch 上。