使用 CUDA 和 NVDEC 在 GPU 上加速影片解碼

TorchCodec 可以利用受支援的 Nvidia 硬體（請參閱支援矩陣 此處）來加速影片解碼。這被稱為“CUDA 解碼”，它分別使用 Nvidia 的 NVDEC 硬體解碼器 和 CUDA 核心來解壓縮和轉換為 RGB。對於實際的解碼步驟以及後續的轉換步驟（如縮放、裁剪或旋轉），CUDA 解碼可以比 CPU 解碼更快。這是因為解碼步驟將解碼後的張量保留在 GPU 記憶體中，因此 GPU 在執行轉換步驟之前無需從主記憶體中獲取資料。編碼包通常比解碼幀小得多，因此 CUDA 解碼還可以減少 PCI-e 頻寬的使用。

何時以及何時不使用 CUDA 解碼

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

1. 您正在解碼高解析度影片

2. 您正在解碼大量影片，並且佔用了 CPU 的全部資源

3. 您想在解碼後對解碼後的張量執行全影像轉換，如縮放或卷積

4. 您的 CPU 已飽和，您想釋放它以供其他工作使用

以下是 CUDA 解碼可能不適用的情況

1. 您需要與 CPU 解碼進行位精確結果比較

2. 您正在處理低解析度影片，而 PCI-e 傳輸延遲很高

3. 您的 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')