安装 | 简单示例 | 详细示例 | 文档 | 贡献指南 | 许可证

TorchCodec

TorchCodec 是一个用于将视频和音频数据解码为 PyTorch 张量的 Python 库，支持 CPU 和 CUDA GPU。它同时支持在 CPU 上进行视频和音频编码！其目标是快速、易用，并与 PyTorch 生态系统良好集成。如果你想使用 PyTorch 在视频和音频上训练机器学习模型，TorchCodec 就是你将其转换为数据的方式。

我们通过以下方式实现这些功能：

• 遵循 Python 和 PyTorch 惯例的 Pythonic API。
• 依赖 FFmpeg 进行解码和编码。TorchCodec 使用你系统上已安装的 FFmpeg 版本。FFmpeg 是一个成熟且覆盖广泛的库，在大多数系统上都可用。然而，它并不易用。TorchCodec 抽象了 FFmpeg 的复杂性，确保其被正确且高效地使用。
• 将数据作为 PyTorch 张量返回，可以直接输入到 PyTorch 转换中或直接用于训练模型。

使用 TorchCodec

以下是 TorchCodec 功能的简要总结。更多详细示例，请查看我们的文档！

解码

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]  # 形状为 [C, H, W] 的 uint8 张量
decoder[0 : -1 : 20]  # 形状为 [N, C, H, W] 的 uint8 堆叠张量

# 索引，附带 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)

片段采样

from torchcodec.samplers import clips_at_regular_timestamps

clips_at_regular_timestamps(
    decoder,
    seconds_between_clip_starts=1.5,
    num_frames_per_clip=4,
    seconds_between_frames=0.1
)
# FrameBatch:
#   data (shape): torch.Size([9, 4, 3, 270, 480])
#   pts_seconds: tensor([[ 0.0000,  0.0667,  0.1668,  0.2669],
#         [ 1.4681,  1.5682,  1.6683,  1.7684],
#         [ 2.9696,  3.0697,  3.1698,  3.2699],
#         ... (输出已截断), dtype=torch.float64)
#   duration_seconds: tensor([[0.0334, 0.0334, 0.0334, 0.0334],
#         [0.0334, 0.0334, 0.0334, 0.0334],
#         [0.0334, 0.0334, 0.0334, 0.0334],
#         ... (输出已截断), dtype=torch.float64)

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

fontfile=/usr/share/fonts/dejavu-sans-mono-fonts/DejaVuSansMono-Bold.ttf
output_video_file=/tmp/output_video.mp4

ffmpeg -f lavfi -i \
    color=size=640x400:duration=10:rate=25:color=blue \
    -vf "drawtext=fontfile=${fontfile}:fontsize=30:fontcolor=white:x=(w-text_w)/2:y=(h-text_h)/2:text='Frame %{frame_num}'" \
    ${output_video_file}

安装 TorchCodec

安装仅支持 CPU 的 TorchCodec

1. 按照官方说明安装最新稳定版本的 PyTorch。对于其他版本，请参考下表了解 torch 和 torchcodec 版本之间的兼容性。

2. 安装 FFmpeg（如果尚未安装）。TorchCodec 支持 [4, 8] 范围内的所有主要 FFmpeg 版本。 Linux 发行版通常预装了 FFmpeg。你需要附带独立共享库的 FFmpeg。这对于 Windows 用户尤其重要：这些版本通常被称为 "shared" 发行版。

   如果尚未安装 FFmpeg，或者你需要更新的版本，一个简单的方法是使用 conda：

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

3. 安装 TorchCodec：

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

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

torchcodec	torch	Python
main / nightly	main / nightly	>=3.10, <=3.14
0.11	2.11	>=3.10, <=3.14
0.10	2.10	>=3.10, <=3.14
0.9	2.9	>=3.10, <=3.14
0.8	2.9	>=3.10, <=3.13
0.7	2.8	>=3.9, <=3.13
0.6	2.8	>=3.9, <=3.13
0.5	2.7	>=3.9, <=3.13
0.4	2.7	>=3.9, <=3.13
0.3	2.7	>=3.9, <=3.13
0.2	2.6	>=3.9, <=3.13
0.1	2.5	>=3.9, <=3.12
0.0.3	2.4	>=3.8, <=3.12

安装支持 CUDA 的 TorchCodec

首先，请确保您的 GPU 具备 NVDEC 硬件，能够解码您所需的格式。更多详情请参考 Nvidia 的 GPU 支持矩阵 此处。

1. 安装支持 NVDEC 的 FFmpeg。 支持 CUDA 的 TorchCodec 应与 FFmpeg 版本 [4, 8] 兼容。

   如果尚未安装 FFmpeg，或者您需要更新的版本，一个简单的安装方法是使用 conda：

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

   安装 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 -i test/resources/nasa_13013.mp4 -f null -
   

Linux

2. 使用 官方说明 安装与您的 CUDA 工具包对应的 PyTorch。您将需要 libnpp 和 libnvrtc CUDA 库，这些通常是 CUDA 工具包的一部分。

3. 安装 TorchCodec

   在 Linux 上，pip install torchcodec 默认安装 CUDA 版本的 wheel 包，这与 pip install torch 的默认行为一致。

   pip install torchcodec
   

   使用 --index-url 来选择不同的 CUDA 工具包版本：

   # 这对应 CUDA 工具包版本 13.0。它应该与您安装 PyTorch 时使用的版本相同（如果您使用 pip 安装 PyTorch）。
   pip install torchcodec --index-url=https://download.pytorch.org/whl/cu130
   

Windows

2. 在 Windows 上（实验性支持），您需要依赖 conda 来安装 PyTorch 和 TorchCodec：

   conda install -c conda-forge "torchcodec=*=*cuda*"
   

基准测试结果

以下结果是在一台负载较轻的 22 核机器上运行 我们的基准测试脚本 生成的，该机器配备了具有 5 个 NVDEC 解码器 的 Nvidia A100 GPU。

顶行是使用 FFmpeg 生成的 曼德博集合 视频，分辨率为 1280x720，60 fps，时长为 120 秒。 底行是 NASA 的宣传视频，分辨率为 960x540，29.7 fps，时长为 206 秒。两个视频均使用 libx264 编码和 yuv420p 像素格式。除 TorchVision 外，所有解码器均使用 FFmpeg 6.1.2。TorchVision 使用 FFmpeg 4.2.2。

对于 TorchCodec，"approx" 标签表示其使用了 近似模式 进行跳转。

贡献

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

许可证

TorchCodec 在 BSD 3 许可证 下发布。

但是，TorchCodec 可能与并非由 Meta 编写的代码一起使用，这些代码可能根据不同的许可证分发。

例如，如果您使用 ENABLE_CUDA=1 构建 TorchCodec 或使用支持 CUDA 的 torchcodec 版本，请在此处查看 CUDA 的许可证： Nvidia 许可证。

TorchCodec 快速上手指南

TorchCodec 是一个用于将视频和音频数据解码为 PyTorch 张量的 Python 库，支持 CPU 和 CUDA GPU。它同样支持在 CPU 上进行视频和音频编码！其目标是快速、易用，并与 PyTorch 生态系统深度集成。如果你希望使用 PyTorch 训练视频和音频的机器学习模型，TorchCodec 就是你处理数据的关键工具。

环境准备

系统要求

• Python: 版本要求取决于你选择的 TorchCodec 版本，请参考下方兼容性表格。
• PyTorch: 需要预先安装与 TorchCodec 版本兼容的 PyTorch。
• FFmpeg: TorchCodec 依赖 FFmpeg（版本 4.x 至 8.x）进行编解码。请确保系统中已安装 FFmpeg 并包含共享库。

版本兼容性

torchcodec 版本	兼容的 torch 版本	兼容的 Python 版本
main / nightly	main / nightly	>=3.10, <=3.14
0.11	2.11	>=3.10, <=3.14
0.10	2.10	>=3.10, <=3.14
0.9	2.9	>=3.10, <=3.14
0.8	2.9	>=3.10, <=3.13
0.7	2.8	>=3.9, <=3.13
0.6	2.8	>=3.9, <=3.13
0.5	2.7	>=3.9, <=3.13
0.4	2.7	>=3.9, <=3.13
0.3	2.7	>=3.9, <=3.13
0.2	2.6	>=3.9, <=3.13
0.1	2.5	>=3.9, <=3.12
0.0.3	2.4	>=3.8, <=3.12

安装步骤

1. 安装 PyTorch

请根据你的系统环境（CPU/CUDA）和 TorchCodec 版本要求，前往 PyTorch 官网 获取安装命令。

2. 安装 FFmpeg

如果你的系统尚未安装 FFmpeg，推荐使用 Conda 安装：

conda install "ffmpeg"
# 或使用 conda-forge 频道
conda install "ffmpeg" -c conda-forge

对于 Windows 用户：请确保安装的是包含共享库（通常标记为 “shared”）的 FFmpeg 版本。

3. 安装 TorchCodec

CPU 版本安装

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

CUDA 版本安装（Linux）

首先，请确认你的 GPU 支持 NVDEC 硬件解码，并已安装对应版本的 CUDA Toolkit。

# 默认安装（通常对应最新的 CUDA 版本）
pip install torchcodec
# 或指定 CUDA Toolkit 版本（例如 CUDA 13.0）
pip install torchcodec --index-url=https://download.pytorch.org/whl/cu130

CUDA 版本安装（Windows，实验性支持）

Windows 用户建议通过 Conda 安装：

conda install -c conda-forge "torchcodec=*=*cuda*"

安装后验证：对于 CUDA 版本，可运行以下命令检查 FFmpeg 的 NVDEC 解码器是否可用：

ffmpeg -decoders | grep -i nvidia
# 预期输出类似： V..... h264_cuvid           Nvidia CUVID H264 decoder (codec h264)

基本使用

以下是一个最简单的视频解码示例，展示如何加载视频并获取帧数据。

from torchcodec.decoders import VideoDecoder

# 初始化解码器，可指定设备为 "cpu" 或 "cuda"
device = "cpu"
decoder = VideoDecoder("path/to/your/video.mp4", device=device)

# 查看视频元数据
print(decoder.metadata)
# 输出示例：
# VideoStreamMetadata:
#   num_frames: 250
#   duration_seconds: 10.0
#   bit_rate: 31315.0
#   codec: h264
#   average_fps: 25.0

# 使用简单的索引 API 获取单帧（返回 uint8 张量，形状为 [C, H, W]）
frame_tensor = decoder[0]

# 使用切片获取多帧（返回 uint8 张量，形状为 [N, C, H, W]）
frames_tensor = decoder[0:-1:20]

# 获取指定索引的帧，并包含时间戳信息
frame_batch = decoder.get_frames_at(indices=[2, 100])
print(frame_batch.data.shape)  # 例如：torch.Size([2, 3, 270, 480])

剪辑采样示例

from torchcodec.samplers import clips_at_regular_timestamps

# 从视频中按固定时间间隔采样剪辑
clips = clips_at_regular_timestamps(
    decoder,
    seconds_between_clip_starts=1.5,  # 剪辑起始点间隔（秒）
    num_frames_per_clip=4,            # 每个剪辑包含的帧数
    seconds_between_frames=0.1        # 剪辑内帧间隔（秒）
)
print(clips.data.shape)  # 例如：torch.Size([9, 4, 3, 270, 480])