From 5d611e7ac2bc3f23c7c694a6e6316167eae6914c Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Wed, 20 Nov 2024 14:24:14 +0000 Subject: [PATCH 1/2] Expose None as option to ffmpeg_num_threads --- src/torchcodec/decoders/_video_decoder.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/torchcodec/decoders/_video_decoder.py b/src/torchcodec/decoders/_video_decoder.py index bd1022de..bf3f9145 100644 --- a/src/torchcodec/decoders/_video_decoder.py +++ b/src/torchcodec/decoders/_video_decoder.py @@ -43,10 +43,11 @@ class VideoDecoder: cheap no-copy operation that allows these frames to be transformed using the `torchvision transforms `_. - num_ffmpeg_threads (int, optional): The number of threads to use for decoding. + num_ffmpeg_threads (int or None, optional): The number of threads to use for decoding. Use 1 for single-threaded decoding which may be best if you are running multiple 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``. + ``None`` is equivalent to passing 0 and lets FFmpeg automatically decide. Default: 1. device (str or torch.device, optional): The device to use for decoding. Default: "cpu". @@ -64,7 +65,7 @@ def __init__( *, stream_index: Optional[int] = None, dimension_order: Literal["NCHW", "NHWC"] = "NCHW", - num_ffmpeg_threads: int = 1, + num_ffmpeg_threads: Optional[int] = 1, device: Optional[Union[str, device]] = "cpu", ): if isinstance(source, str): From ba5a13e7b248808796a4b1d1d823adf474e6a7ea Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Wed, 20 Nov 2024 14:27:05 +0000 Subject: [PATCH 2/2] Accept 'auto' in num_ffmpeg_threads --- src/torchcodec/decoders/_video_decoder.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/src/torchcodec/decoders/_video_decoder.py b/src/torchcodec/decoders/_video_decoder.py index bf3f9145..0494feec 100644 --- a/src/torchcodec/decoders/_video_decoder.py +++ b/src/torchcodec/decoders/_video_decoder.py @@ -43,11 +43,11 @@ class VideoDecoder: cheap no-copy operation that allows these frames to be transformed using the `torchvision transforms `_. - num_ffmpeg_threads (int or None, optional): The number of threads to use for decoding. + num_ffmpeg_threads (int or str, optional): The number of threads to use for decoding. Use 1 for single-threaded decoding which may be best if you are running multiple 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``. - ``None`` is equivalent to passing 0 and lets FFmpeg automatically decide. + "auto" is equivalent to passing 0 and lets FFmpeg automatically decide. Default: 1. device (str or torch.device, optional): The device to use for decoding. Default: "cpu". @@ -65,7 +65,7 @@ def __init__( *, stream_index: Optional[int] = None, dimension_order: Literal["NCHW", "NHWC"] = "NCHW", - num_ffmpeg_threads: Optional[int] = 1, + num_ffmpeg_threads: Union[int, str] = 1, device: Optional[Union[str, device]] = "cpu", ): if isinstance(source, str): @@ -89,6 +89,14 @@ def __init__( f"Supported values are {', '.join(allowed_dimension_orders)}." ) + if isinstance(num_ffmpeg_threads, str): + if num_ffmpeg_threads != "auto": + raise ValueError( + "num_ffmpeg_threads should be 'auto' when it's a string. " + f"Got {num_ffmpeg_threads}" + ) + num_ffmpeg_threads = 0 + core.scan_all_streams_to_update_metadata(self._decoder) core.add_video_stream( self._decoder,