使用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.5.1+cu124'
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=ac51ece7-7469-b9f9-89f5-273763cb89de, 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 0x7f304bf38550>)

使用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 分钟 7.063 秒)

Gallery generated by Sphinx-Gallery