PyTorch 媒体处理生态的演进 – PyTorch

来源：https://pytorch.org/blog/evolution-of-the-pytorch-media-processing-landscape/

重点项目

• PyTorch logo (https://pytorch.org/projects/pytorch/)

太长不看（TL;DR）

如果你需要解码或编码媒体（无论是图像、视频还是音频），请使用 TorchCodec (https://github.com/pytorch/torchcodec)。如果你需要变换媒体，请对图像和视频使用 TorchVision (https://github.com/pytorch/vision)，对音频使用 TorchAudio (https://github.com/pytorch/audio)。这些库中原本包含的其他内容（模型、数据集等），如今由更广泛的生态提供更好的支持，例如 HuggingFace 的各个库。

图像、视频和音频如今已成为大量模型开发的核心——从视觉-语言模型到生成图像和视频的扩散模型皆是如此。训练这些模型意味着将媒体解码为张量并进行变换，而生成模型还需要将其输出重新编码回媒体文件。过去两年中，我们将 PyTorch 的媒体技术栈整合到了三个库中：TorchCodec (https://github.com/pytorch/torchcodec)、TorchVision (https://github.com/pytorch/vision) 和 TorchAudio (https://github.com/pytorch/audio)。本文介绍最终的架构布局以及背后的原因。

过去与现在

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 Codec SDK，以及每种图像格式一个库（libjpeg、libpng 等），每个库在我们能够发布什么内容方面都有各自的许可规则。除此之外，它们全都不是纯 Python 依赖——都是 C/C++ 库，这使得构建、CI 和发布都困难得多。如今所有这些复杂性都集中在 TorchCodec 中，这极大地简化了 TorchVision 和 TorchAudio 的维护负担。

TorchVision 和 TorchAudio：如今专注于媒体变换

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

背后的驱动因素在于我们观察到：社区中最活跃的使用场景正是 TorchVision 和 TorchAudio 中的变换功能。相比之下，模型、数据集和流水线的使用相对落后，因为已经涌现出优秀的替代方案，例如 HuggingFace 的各个库。缩小 TorchAudio 和 TorchVision 的范围，让我们能够集中投入其优势领域，以及生态系统其他部分未覆盖的能力。

这次转型对 TorchAudio 尤其具有颠覆性：大量 API 被弃用并最终移除。我们感谢社区的耐心与反馈，这让我们得以重新审视最初的方案。我们原本计划移除的几个流行 API 最终得以保留（https://github.com/pytorch/audio/issues/3902#issuecomment-3779752106）。这次迁移很艰难，但我们认为是值得的。更小的 TorchAudio 才是真正可维护的 TorchAudio，也才是用户可以持续信赖的 TorchAudio。

三者如何协同工作

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

对于视频，结合 TorchCodec 的 VideoDecoder (https://meta-pytorch.org/torchcodec/main/generated/torchcodec.decoders.VideoDecoder.html) 和 TorchVision 的 v2 变换 (https://docs.pytorch.org/vision/stable/transforms.html)：

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 (https://meta-pytorch.org/torchcodec/main/generated/torchcodec.decoders.AudioDecoder.html) 和 TorchAudio 的变换 (https://docs.pytorch.org/audio/stable/transforms.html)：

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 (https://meta-pytorch.org/torchcodec/main/generated/torchcodec.decoders.decode_image.html#torchcodec.decoders.decode_image) 及相关函数返回普通张量，可直接输入 torchvision.transforms.v2。编码则走相反方向，由 VideoEncoder (https://meta-pytorch.org/torchcodec/main/generated/torchcodec.encoders.VideoEncoder.html)、AudioEncoder (https://meta-pytorch.org/torchcodec/main/generated/torchcodec.encoders.AudioEncoder.html)、JpegEncoder (https://meta-pytorch.org/torchcodec/main/generated/torchcodec.encoders.JpegEncoder.html) 和 PngEncoder (https://meta-pytorch.org/torchcodec/main/generated/torchcodec.encoders.PngEncoder.html) 完成。

在发布方面，这三个库如今均已实现 ABI 稳定 (https://docs.pytorch.org/docs/2.14/notes/libtorch_stable_abi.html)。这意味着某个版本的 TorchCodec、TorchVision 或 TorchAudio 不再绑定到单一的 PyTorch 版本：它可以与之后发布的 PyTorch 版本继续协同工作。它们不再需要为每个 PyTorch 发布版本重新构建，因此其发布节奏也不再与 PyTorch 保持一致。这是完全正常的！

如果你仍在使用 TorchVision 或 TorchAudio 中的解码或编码 API，现在正是迁移到 TorchCodec 的好时机：可以从我们的迁移指南 (https://meta-pytorch.org/torchcodec/main/generated_examples/migration/torchvision_migration.html) 开始，如果你需要的功能有所缺失，欢迎在 GitHub (https://github.com/meta-pytorch/torchcodec/issues) 上告诉我们。

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