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

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

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

在某些情况下，CUDA 解码可以提供比 CPU 解码更快的速度

您正在解码高分辨率视频
您正在解码大量视频，并且占用了 CPU 的全部资源
您想在解码后对解码后的张量执行全图像转换，如缩放或卷积
您的 CPU 已饱和，您想释放它以供其他工作使用

以下是 CUDA 解码可能不适用的情况

您需要与 CPU 解码进行位精确结果比较
您正在处理低分辨率视频，而 PCI-e 传输延迟很高
您的 GPU 已经很忙，而 CPU 不是很忙

最好尝试 CUDA 解码，看看它是否能改进您的用例。使用 TorchCodec，您只需将设备参数传递给 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.9.0.dev20250714+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=1bb3fb37-9a86-0064-473b-947e3a43d91c, pci_bus_id=0, pci_device_id=30, pci_domain_id=0, L2_cache_size=2MB)

下载视频¶

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

编解码器：H.264
分辨率：960x540
帧率：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 0x7fa45c833970>)

使用 VideoDecoder 进行 CUDA 解码¶

要使用 CUDA 解码器，您需要将 cuda 设备传递给解码器。

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 数学不是位精确的。

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.664 秒)

由 Sphinx-Gallery 生成的画廊