meta-pytorch/torchcodec 来源：https://github.com/meta-pytorch/torchcodec 安装 | 简单示例 | 详细示例 (https://meta-pytorch.org/torchcodec/stable/generated_examples/) | 文档 (https://meta-pytorch.org/torchcodec) | 贡献指南 | 许可协议

TorchCodec

TorchCodec 是一个 Python 库，用于将视频和音频数据解码为 PyTorch 张量，支持 CPU 和 CUDA GPU。它还支持 CPU 上的视频和音频编码！它的目标是快速、易用，并与 PyTorch 生态系统良好集成。

如果你想要使用 PyTorch 训练视频和音频的机器学习模型，或者进行推理，TorchCodec 就是将这些数据转化为张量的工具。我们通过以下能力实现这些目标：

• 遵循 Python 和 PyTorch 惯例的 Pythonic API。
• 依赖 FFmpeg (https://www.ffmpeg.org/) 进行解码和编码。TorchCodec 使用你已经安装的 FFmpeg 版本。FFmpeg 是一个成熟的库，覆盖广泛，在大多数系统上可用。但它使用起来并不容易。TorchCodec 抽象了 FFmpeg 的复杂性，确保其正确高效地被使用。
• 以 PyTorch 张量的形式返回数据，可直接用于 PyTorch 变换或直接训练模型。

使用示例

以下是一些使用 TorchCodec 的示例。更多详细示例和用例，请查看我们的文档 (https://meta-pytorch.org/torchcodec/stable/generated_examples/)！

视频解码

from torchcodec.decoders import VideoDecoder

device = "cpu" # 或例如 "cuda" 
decoder = VideoDecoder("path/to/video.mp4", device=device)

decoder.metadata
# VideoStreamMetadata:
#   num_frames: 250
#   duration_seconds: 10.0
#   bit_rate: 31315.0
#   codec: h264
#   average_fps: 25.0
#   ... (输出已截断)

# 简单的索引 API
decoder[0]          # uint8 张量，形状 [C, H, W]
decoder[0 : -1 : 20] # uint8 堆叠张量，形状 [N, C, H, W]

# 索引，带 PTS 和持续时间信息：
decoder.get_frames_at(indices=[2, 100])
# FrameBatch:
#   data (shape): torch.Size([2, 3, 270, 480])
#   pts_seconds: tensor([0.0667, 3.3367], dtype=torch.float64)
#   duration_seconds: tensor([0.0334, 0.0334], dtype=torch.float64)

# 基于时间的索引，带 PTS 和持续时间信息
decoder.get_frames_played_at(seconds=[0.5, 10.4])
# FrameBatch:
#   data (shape): torch.Size([2, 3, 270, 480])
#   pts_seconds: tensor([ 0.4671, 10.3770], dtype=torch.float64)
#   duration_seconds: tensor([0.0334, 0.0334], dtype=torch.float64)

你可以使用以下代码片段通过 FFmpeg 生成一个视频，并尝试使用 VideoDecoder：

ffmpeg -f lavfi -i testsrc2=size=640x400:duration=10:rate=25 /tmp/output_video.mp4

编码

from torchcodec.encoders import Encoder

encoder = Encoder()
video_stream = encoder.add_video(
    height=height,
    width=width,
    frame_rate=frame_rate,
)
audio_stream = encoder.add_audio(
    sample_rate=sample_rate,
    num_channels=num_channels,
)

with encoder.open_file("output.mp4"):
    video_stream.add_frames(frames_batch_0)
    audio_stream.add_samples(samples_batch_0)
    video_stream.add_frames(frames_batch_1)
    audio_stream.add_samples(samples_batch_1)
    # ...

安装 TorchCodec

1. 安装 FFmpeg（如果尚未安装）。TorchCodec 支持 [4, 8] 范围内的所有主要 FFmpeg 版本。Linux 发行版通常预装了 FFmpeg。你需要使用带有单独共享库的 FFmpeg。这对 Windows 用户尤其重要：这些通常被称为“shared”版本。如果 FFmpeg 尚未安装，或者你需要更新的版本，一种简单的安装方法是使用 conda：

   conda install "ffmpeg"
   # 或者
   conda install "ffmpeg" -c conda-forge
   

2. 安装 PyTorch 和 TorchCodec：

   pip install torch torchcodec
   

就是这样！在 Linux x86 和 aarch64 上，这默认会安装支持 CUDA 的轮子（与 pip install torch 的默认行为一致）。即使你的机器上没有 GPU，这些轮子仍然应该能工作。在 macOS 和 Windows 上，这会安装仅 CPU 的轮子。CPU 轮子可在 Linux（x86_64 和 aarch64）、macOS 和 Windows 上使用。

CUDA 支持

在 Linux 上默认安装支持 CUDA 的轮子。对于 Windows，你需要像下面描述的那样传递 --index-url。确保你有一块带有 NVDEC 硬件并能解码所需格式的 GPU。参考 Nvidia 的 GPU 支持矩阵 (https://developer.nvidia.com/video-encode-and-decode-gpu-support-matrix-new)。你需要 libnvrtc CUDA 库，它通常是 CUDA Toolkit 的一部分。

要选择特定的 CUDA Toolkit 版本，请使用 --index-url。确保也安装对应的 PyTorch 版本（参考官方安装指南 (https://pytorch.org/get-started/locally/)）：

# 这对应于 CUDA Toolkit 版本 13.0。
pip install torch torchcodec --index-url=https://download.pytorch.org/whl/cu130

确保你的 FFmpeg 支持 NVDEC：

ffmpeg -decoders | grep -i nvidia
# 应该显示类似下面的一行：
# V..... h264_cuvid Nvidia CUVID H264 decoder (codec h264)

要检查 FFmpeg 库是否与 NVDEC 正确工作，你可以解码一个生成的测试视频：

ffmpeg -hwaccel cuda -hwaccel_output_format cuda -f lavfi -i testsrc2=duration=1 -f null -

仅 CPU 安装

要显式安装仅 CPU 的轮子（例如在 Linux 上，CUDA 轮子是默认的）：

pip install torchcodec --index-url=https://download.pytorch.org/whl/cpu

XPU 支持

Intel GPU（XPU）支持需要 TorchCodec 的独立插件：

pip install torchcodec-xpu --extra-index-url=https://download.pytorch.org/whl/xpu

任何与 XPU 相关的支持，请参考 https://github.com/intel/torchlib-xpu。

与 torch 版本的兼容性

下表显示了 torchcodec、torch 和 Python 版本之间的兼容性。

贡献

我们欢迎对 TorchCodec 的贡献！更多详情请参阅我们的贡献指南。

许可协议

TorchCodec 采用 BSD 3 许可协议 发布。然而，TorchCodec 可能与非 Meta 编写的代码一起使用，这些代码可能采用不同的许可协议。例如，如果你使用 ENABLE_CUDA=1 构建 TorchCodec 或使用支持 CUDA 的 torchcodec 版本，请在此处查看 CUDA 的许可协议：Nvidia 许可 (https://docs.nvidia.com/cuda/eula/index.html)。