.

Author: ahmadsharif1

benchmarks/samplers/benchmark_samplers.py

@@ -1,3 +1,4 @@
+import argparse
 from pathlib import Path
 from time import perf_counter_ns
 
@@ -45,51 +46,72 @@ def report_stats(times, num_frames, unit="ms"):
     return med, fps
 
 
-def sample(sampler, **kwargs):
-    decoder = VideoDecoder(VIDEO_PATH)
+def sample(decoder, sampler, **kwargs):
     return sampler(
         decoder,
         num_frames_per_clip=10,
         **kwargs,
     )
 
 
-VIDEO_PATH = Path(__file__).parent / "../../test/resources/nasa_13013.mp4"
-NUM_EXP = 30
-
-for num_clips in (1, 50):
-    print("-" * 10)
-    print(f"{num_clips = }")
-
-    print("clips_at_random_indices     ", end="")
-    times, num_frames = bench(
-        sample, clips_at_random_indices, num_clips=num_clips, num_exp=NUM_EXP, warmup=2
-    )
-    report_stats(times, num_frames, unit="ms")
-
-    print("clips_at_regular_indices    ", end="")
-    times, num_frames = bench(
-        sample, clips_at_regular_indices, num_clips=num_clips, num_exp=NUM_EXP, warmup=2
-    )
-    report_stats(times, num_frames, unit="ms")
-
-    print("clips_at_random_timestamps  ", end="")
-    times, num_frames = bench(
-        sample,
-        clips_at_random_timestamps,
-        num_clips=num_clips,
-        num_exp=NUM_EXP,
-        warmup=2,
-    )
-    report_stats(times, num_frames, unit="ms")
-
-    print("clips_at_regular_timestamps ", end="")
-    seconds_between_clip_starts = 13 / num_clips  # approximate. video is 13s long
-    times, num_frames = bench(
-        sample,
-        clips_at_regular_timestamps,
-        seconds_between_clip_starts=seconds_between_clip_starts,
-        num_exp=NUM_EXP,
-        warmup=2,
-    )
-    report_stats(times, num_frames, unit="ms")
+def main(device, video):
+    NUM_EXP = 30
+
+    for num_clips in (1, 50):
+        print("-" * 10)
+        print(f"{num_clips = }")
+
+        print("clips_at_random_indices     ", end="")
+        decoder = VideoDecoder(video, device=device)
+        times, num_frames = bench(
+            sample,
+            decoder,
+            clips_at_random_indices,
+            num_clips=num_clips,
+            num_exp=NUM_EXP,
+            warmup=2,
+        )
+        report_stats(times, num_frames, unit="ms")
+
+        print("clips_at_regular_indices    ", end="")
+        times, num_frames = bench(
+            sample,
+            decoder,
+            clips_at_regular_indices,
+            num_clips=num_clips,
+            num_exp=NUM_EXP,
+            warmup=2,
+        )
+        report_stats(times, num_frames, unit="ms")
+
+        print("clips_at_random_timestamps  ", end="")
+        times, num_frames = bench(
+            sample,
+            decoder,
+            clips_at_random_timestamps,
+            num_clips=num_clips,
+            num_exp=NUM_EXP,
+            warmup=2,
+        )
+        report_stats(times, num_frames, unit="ms")
+
+        print("clips_at_regular_timestamps ", end="")
+        seconds_between_clip_starts = 13 / num_clips  # approximate. video is 13s long
+        times, num_frames = bench(
+            sample,
+            decoder,
+            clips_at_regular_timestamps,
+            seconds_between_clip_starts=seconds_between_clip_starts,
+            num_exp=NUM_EXP,
+            warmup=2,
+        )
+        report_stats(times, num_frames, unit="ms")
+
+
+if __name__ == "__main__":
+    DEFAULT_VIDEO_PATH = Path(__file__).parent / "../../test/resources/nasa_13013.mp4"
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", type=str, default="cpu")
+    parser.add_argument("--video", type=str, default=str(DEFAULT_VIDEO_PATH))
+    args = parser.parse_args()
+    main(args.device, args.video)

src/torchcodec/decoders/_core/CPUOnlyDevice.cpp

@@ -19,7 +19,8 @@ void convertAVFrameToDecodedOutputOnCuda(
     const VideoDecoder::VideoStreamDecoderOptions& options,
     AVCodecContext* codecContext,
     VideoDecoder::RawDecodedOutput& rawOutput,
-    VideoDecoder::DecodedOutput& output) {
+    VideoDecoder::DecodedOutput& output,
+    std::optional<torch::Tensor> preAllocatedOutputTensor) {
   throwUnsupportedDeviceError(device);
 }
 

src/torchcodec/decoders/_video_decoder.py

@@ -8,7 +8,7 @@
 from pathlib import Path
 from typing import Literal, Optional, Tuple, Union
 
-from torch import Tensor
+from torch import device, Tensor
 
 from torchcodec import Frame, FrameBatch
 from torchcodec.decoders import _core as core
@@ -41,6 +41,8 @@ class VideoDecoder:
             instances of ``VideoDecoder`` in parallel. Use a higher number for multi-threaded
             decoding which is best if you are running a single instance of ``VideoDecoder``.
             Default: 1.
+        device (str or torch.device, optional): The device to use for decoding.
+
 
             .. note::
 
@@ -64,6 +66,7 @@ def __init__(
         stream_index: Optional[int] = None,
         dimension_order: Literal["NCHW", "NHWC"] = "NCHW",
         num_ffmpeg_threads: int = 1,
+        device: Optional[Union[str, device]] = "cpu",
     ):
         if isinstance(source, str):
             self._decoder = core.create_from_file(source)
@@ -92,6 +95,7 @@ def __init__(
             stream_index=stream_index,
             dimension_order=dimension_order,
             num_threads=num_ffmpeg_threads,
+            device=device,
         )
 
         self.metadata, self.stream_index = _get_and_validate_stream_metadata(