在 GPU 上使用 CUDA 和 NVDEC 加速视频解码¶

TorchCodec 可以使用支持的 Nvidia 硬件（参见此处的支持矩阵）来加速视频解码。这称为“CUDA 解码”，它分别使用 Nvidia 的 NVDEC 硬件解码器和 CUDA kernel 来进行解压缩并转换为 RGB。对于实际的解码步骤以及后续的转换步骤（如缩放、裁剪或旋转），CUDA 解码可能比 CPU 解码更快。这是因为解码步骤将解码后的张量留在 GPU 内存中，因此 GPU 在运行转换步骤之前无需从主内存中获取数据。编码后的数据包通常比解码后的帧小得多，因此 CUDA 解码也使用更少的 PCI-e 带宽。

何时使用 CUDA 解码及何时不使用¶

在以下几种场景下，CUDA 解码可以比 CPU 解码提供速度提升

你正在解码高分辨率视频
你正在解码大量视频，导致 CPU 饱和
解码后你想对解码后的张量进行整图转换，如缩放或卷积
你的 CPU 已经饱和，你想将其用于其他工作

在以下情况下，CUDA 解码可能不适用

你需要与 CPU 解码完全一致（bit-exact）的结果
你的视频分辨率较低且 PCI-e 传输延迟较大
你的 GPU 已经很忙碌而 CPU 并不忙

最好通过实验来确定 CUDA 解码是否能改善你的使用场景。使用 TorchCodec，你只需将 device 参数传递给 VideoDecoder 类即可使用 CUDA 解码。

安装启用 CUDA 的 TorchCodec¶

请参考 README 中的安装指南。

检查 PyTorch 是否已启用 CUDA¶

注意

本教程需要编译时启用 CUDA 支持的 FFmpeg 库。

import torch

print(f"{torch.__version__=}")
print(f"{torch.cuda.is_available()=}")
print(f"{torch.cuda.get_device_properties(0)=}")

torch.__version__='2.8.0.dev20250423+cu126'
torch.cuda.is_available()=True
torch.cuda.get_device_properties(0)=_CudaDeviceProperties(name='Tesla M60', major=5, minor=2, total_memory=7606MB, multi_processor_count=16, uuid=c27ee195-a507-bd7c-643e-0888b2168fa5, pci_bus_id=0, pci_device_id=30, pci_domain_id=0, L2_cache_size=2MB)

下载视频¶

我们将使用以下视频，它具有以下属性

编码格式: H.264
分辨率: 960x540
帧率 (FPS): 29.97
像素格式: YUV420P

import urllib.request

video_file = "video.mp4"
urllib.request.urlretrieve(
    "https://download.pytorch.org/torchaudio/tutorial-assets/stream-api/NASAs_Most_Scientifically_Complex_Space_Observatory_Requires_Precision-MP4_small.mp4",
    video_file,
)

('video.mp4', <http.client.HTTPMessage object at 0x7fd1a04cf4f0>)

使用 VideoDecoder 进行 CUDA 解码¶

要使用 CUDA 解码器，你需要向解码器传入一个 CUDA device。

from torchcodec.decoders import VideoDecoder

decoder = VideoDecoder(video_file, device="cuda")
frame = decoder[0]

视频帧被解码并以 NCHW 格式的张量返回。

print(frame.shape, frame.dtype)

torch.Size([3, 540, 960]) torch.uint8

视频帧保留在 GPU 内存中。

print(frame.data.device)

cuda:0

可视化帧¶

让我们看看 CUDA 解码器解码的帧，并将其与 CPU 解码器的等效结果进行比较。

timestamps = [12, 19, 45, 131, 180]
cpu_decoder = VideoDecoder(video_file, device="cpu")
cuda_decoder = VideoDecoder(video_file, device="cuda")
cpu_frames = cpu_decoder.get_frames_played_at(timestamps).data
cuda_frames = cuda_decoder.get_frames_played_at(timestamps).data


def plot_cpu_and_cuda_frames(cpu_frames: torch.Tensor, cuda_frames: torch.Tensor):
    try:
        import matplotlib.pyplot as plt
        from torchvision.transforms.v2.functional import to_pil_image
    except ImportError:
        print("Cannot plot, please run `pip install torchvision matplotlib`")
        return
    n_rows = len(timestamps)
    fig, axes = plt.subplots(n_rows, 2, figsize=[12.8, 16.0])
    for i in range(n_rows):
        axes[i][0].imshow(to_pil_image(cpu_frames[i].to("cpu")))
        axes[i][1].imshow(to_pil_image(cuda_frames[i].to("cpu")))

    axes[0][0].set_title("CPU decoder", fontsize=24)
    axes[0][1].set_title("CUDA decoder", fontsize=24)
    plt.setp(axes, xticks=[], yticks=[])
    plt.tight_layout()


plot_cpu_and_cuda_frames(cpu_frames, cuda_frames)

它们在视觉上看起来与人眼相似，但可能存在细微差异，因为 CUDA 的计算与 CPU 的计算在位级别上不完全一致（bit-exact）。

frames_equal = torch.equal(cpu_frames.to("cuda"), cuda_frames)
mean_abs_diff = torch.mean(
    torch.abs(cpu_frames.float().to("cuda") - cuda_frames.float())
)
max_abs_diff = torch.max(torch.abs(cpu_frames.to("cuda").float() - cuda_frames.float()))
print(f"{frames_equal=}")
print(f"{mean_abs_diff=}")
print(f"{max_abs_diff=}")

frames_equal=False
mean_abs_diff=tensor(0.5636, device='cuda:0')
max_abs_diff=tensor(2., device='cuda:0')

脚本总运行时间： (0 分钟 6.767 秒)

由 Sphinx-Gallery 生成的图库