Streaming nvdec decoding

This example shows how to decode a video with GPU in streaming fashion.

Source

Source

Click here to see the source.
  1# Copyright (c) Meta Platforms, Inc. and affiliates.
  2# All rights reserved.
  3#
  4# This source code is licensed under the BSD-style license found in the
  5# LICENSE file in the root directory of this source tree.
  6
  7"""This example shows how to decode a video with GPU in streaming fashion."""
  8
  9__all__ = [
 10    "main",
 11    "parse_args",
 12    "run",
 13    "decode",
 14    "torch_cuda_warmup",
 15]
 16
 17import argparse
 18import contextlib
 19import logging
 20import pathlib
 21import time
 22
 23import spdl.io
 24import torch
 25from PIL import Image
 26from spdl.io import CUDAConfig
 27from torch.profiler import profile
 28
 29
 30def parse_args(args: list[str] | None = None) -> tuple[argparse.Namespace, list[str]]:
 31    """Parse command line arguments.
 32
 33    Args:
 34        args: The command line arguments. By default it reads ``sys.argv``.
 35
 36    Returns:
 37        Tuple of parsed arguments and unused arguments, as returned by
 38        :py:meth:`argparse.ArgumentParser.parse_known_args`.
 39    """
 40
 41    parser = argparse.ArgumentParser(
 42        description=__doc__,
 43    )
 44    parser.add_argument(
 45        "--input-file", required=True, help="The input video to process."
 46    )
 47    parser.add_argument(
 48        "--plot-dir",
 49        type=pathlib.Path,
 50        help="If provided, plot the result to the given directory.",
 51    )
 52    parser.add_argument(
 53        "--trace-path",
 54        help="If provided, trace the execution. e.g. 'trace.json.gz'",
 55    )
 56    parser.add_argument(
 57        "--device-index",
 58        type=int,
 59        help="The CUDA device index. By default it use the last one.",
 60    )
 61    parser.add_argument(
 62        "--width",
 63        type=int,
 64        default=320,
 65        help="Rescale the video to this width. Provide -1 to disable.",
 66    )
 67    parser.add_argument(
 68        "--height",
 69        type=int,
 70        default=240,
 71        help="Rescale the video to this height. Provide -1 to disable.",
 72    )
 73    return parser.parse_known_args(args)
 74
 75
 76def decode(
 77    src: str,
 78    device_config: CUDAConfig,
 79    post_processing_params: dict[str, int],
 80    profiler: torch.profiler.profile | None,
 81    plot_dir: pathlib.Path | None,
 82) -> None:
 83    """Decode video in streaming fashion with optional resizing, profiling and exporting.
 84
 85    Args:
 86        src: The path or URL to the source video.
 87        device_config: The GPU configuration.
 88        post_processing_params: Post processing argument.
 89            See :py:func:`spdl.io.streaming_load_video_nvdec`.
 90        profiler: PyTorch Profiler or ``None``.
 91        plot_dir: If provided, the decoded frames are exported as images to the directory.
 92    """
 93    streamer = spdl.io.streaming_load_video_nvdec(
 94        src,
 95        device_config,
 96        num_frames=32,
 97        post_processing_params=post_processing_params,
 98    )
 99
100    i, num_frames = 0, 0
101    t0 = time.monotonic()
102    for buffers in streamer:
103        buffer = spdl.io.nv12_to_rgb(buffers, device_config=device_config, sync=True)
104        tensor = spdl.io.to_torch(buffer)
105        num_frames += len(tensor)
106
107        if plot_dir is not None:
108            for f in tensor.permute(0, 2, 3, 1):
109                img = Image.fromarray(f.cpu().numpy())
110                img.save(plot_dir / f"{i:05d}.png")
111                i += 1
112
113        if profiler is not None:
114            profiler.step()
115            if num_frames >= 500:
116                break
117
118    elapsed = time.monotonic() - t0
119    qps = num_frames / elapsed
120    print(f"Processed {num_frames} frames in {elapsed:.1f} sec. QPS: {qps:.1f}")
121
122
123def torch_cuda_warmup(device_index: int | None) -> tuple[int, torch.cuda.Stream]:
124    """Initialize the CUDA context perform dry-run.
125
126    Args:
127        device_index: The CUDA device to use. If ``None``, the last available device is used.
128    """
129    assert torch.cuda.is_available()
130
131    cuda_index: int = device_index or (torch.cuda.device_count() - 1)
132    stream = torch.cuda.Stream(device=cuda_index)
133    with torch.cuda.stream(stream):
134        a = torch.empty([32, 3, 1080, 1920])
135        a.pin_memory().to(f"cuda:{cuda_index}", non_blocking=True)
136    stream.synchronize()
137    return cuda_index, stream
138
139
140def run(
141    src: str,
142    device_index: int | None,
143    post_processing_params: dict[str, int],
144    profiler: torch.profiler.profile,
145    plot_dir: pathlib.Path,
146) -> None:
147    """Run the benchmark."""
148    cuda_index, stream = torch_cuda_warmup(device_index)
149
150    device_config = spdl.io.cuda_config(
151        device_index=cuda_index,
152        allocator=(
153            torch.cuda.caching_allocator_alloc,
154            torch.cuda.caching_allocator_delete,
155        ),
156        stream=stream.cuda_stream,
157    )
158
159    for i in range(3):
160        with torch.autograd.profiler.record_function(f"decode_{i}"):
161            decode(src, device_config, post_processing_params, profiler, plot_dir)
162
163
164def main(args: list[str] | None = None) -> None:
165    """The main entrypoint for the CLI."""
166    ns, _ = parse_args(args)
167
168    logging.basicConfig(level=logging.INFO)
169
170    prof = None
171    post_process = {
172        "scale_width": ns.width if ns.width > 0 else None,
173        "scale_height": ns.height if ns.height > 0 else None,
174    }
175    with contextlib.ExitStack() as stack:
176        if ns.trace_path:
177            prof = stack.enter_context(
178                profile(
179                    with_stack=True,
180                    on_trace_ready=lambda p: p.export_chrome_trace(ns.trace_path),
181                )
182            )
183
184        # pyrefly: ignore [bad-argument-type]
185        run(ns.input_file, ns.device_index, post_process, prof, ns.plot_dir)
186
187
188if __name__ == "__main__":
189    main()

API Reference

Functions

main(args: list[str] | None = None) → None[source]

The main entrypoint for the CLI.

parse_args(args: list[str] | None = None) → tuple[Namespace, list[str]][source]

Parse command line arguments.

Parameters:

args – The command line arguments. By default it reads sys.argv.

Returns:

Tuple of parsed arguments and unused arguments, as returned by argparse.ArgumentParser.parse_known_args().

run(src: str, device_index: int | None, post_processing_params: dict[str, int], profiler: profile, plot_dir: Path) → None[source]

Run the benchmark.

decode(src: str, device_config: CUDAConfig, post_processing_params: dict[str, int], profiler: profile | None, plot_dir: Path | None) → None[source]

Decode video in streaming fashion with optional resizing, profiling and exporting.

Parameters:

  • src – The path or URL to the source video.

  • device_config – The GPU configuration.

  • post_processing_params – Post processing argument. See spdl.io.streaming_load_video_nvdec().

  • profiler – PyTorch Profiler or None.

  • plot_dir – If provided, the decoded frames are exported as images to the directory.

torch_cuda_warmup(device_index: int | None) → tuple[int, Stream][source]

Initialize the CUDA context perform dry-run.

Parameters:

device_index – The CUDA device to use. If None, the last available device is used.