Skip to content

Commit

Permalink
Merge branch 'main' of github.com:pytorch/torchcodec into cuda_agnostic
Browse files Browse the repository at this point in the history
  • Loading branch information
@scotts
scotts committed Nov 21, 2024
2 parents 7bf0a15 + 7e2da5d commit 1550a37
Show file tree
Hide file tree
Showing 3 changed files with 457 additions and 408 deletions.
77 changes: 64 additions & 13 deletions benchmarks/decoders/benchmark_decoders_library.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,7 @@ def get_consecutive_frames_from_video(self, video_file, numFramesToDecode):

def decode_and_transform(self, video_file, pts_list, height, width, device):
self.torchvision.set_video_backend(self._backend)
reader = self.torchvision.io.VideoReader(video_file, "video")
reader = self.torchvision.io.VideoReader(video_file, "video", num_threads=1)
frames = []
for pts in pts_list:
reader.seek(pts)
Expand Down Expand Up @@ -177,15 +177,20 @@ def get_consecutive_frames_from_video(self, video_file, numFramesToDecode):

class TorchCodecCoreNonBatch(AbstractDecoder):
def __init__(self, num_threads=None, color_conversion_library=None, device="cpu"):
self._num_threads = int(num_threads) if num_threads else None
self._num_threads = num_threads
self._color_conversion_library = color_conversion_library
self._device = device

from torchvision.transforms import v2 as transforms_v2

self.transforms_v2 = transforms_v2

def get_frames_from_video(self, video_file, pts_list):
decoder = create_from_file(video_file)
num_threads = int(self._num_threads) if self._num_threads else 0
_add_video_stream(
decoder,
num_threads=self._num_threads,
num_threads=num_threads,
color_conversion_library=self._color_conversion_library,
device=self._device,
)
Expand All @@ -199,10 +204,11 @@ def get_frames_from_video(self, video_file, pts_list):
return frames

def get_consecutive_frames_from_video(self, video_file, numFramesToDecode):
num_threads = int(self._num_threads) if self._num_threads else 0
decoder = create_from_file(video_file)
_add_video_stream(
decoder,
num_threads=self._num_threads,
num_threads=num_threads,
color_conversion_library=self._color_conversion_library,
device=self._device,
)
Expand All @@ -214,6 +220,29 @@ def get_consecutive_frames_from_video(self, video_file, numFramesToDecode):

return frames

def decode_and_transform(self, video_file, pts_list, height, width, device):
num_threads = int(self._num_threads) if self._num_threads else 1
decoder = create_from_file(video_file)
_add_video_stream(
decoder,
num_threads=num_threads,
color_conversion_library=self._color_conversion_library,
device=self._device,
)

frames = []
for pts in pts_list:
seek_to_pts(decoder, pts)
frame, *_ = get_next_frame(decoder)
frames.append(frame)

frames = [
self.transforms_v2.functional.resize(frame.to(device), (height, width))
for frame in frames
]

return frames


class TorchCodecCoreBatch(AbstractDecoder):
def __init__(self, num_threads=None, color_conversion_library=None, device="cpu"):
Expand Down Expand Up @@ -258,22 +287,28 @@ def get_consecutive_frames_from_video(self, video_file, numFramesToDecode):

class TorchCodecPublic(AbstractDecoder):
def __init__(self, num_ffmpeg_threads=None, device="cpu"):
self._num_ffmpeg_threads = int(num_ffmpeg_threads) if num_ffmpeg_threads else 1
self._num_ffmpeg_threads = num_ffmpeg_threads
self._device = device

from torchvision.transforms import v2 as transforms_v2

self.transforms_v2 = transforms_v2

def get_frames_from_video(self, video_file, pts_list):
num_ffmpeg_threads = (
int(self._num_ffmpeg_threads) if self._num_ffmpeg_threads else 0
)
decoder = VideoDecoder(
video_file, num_ffmpeg_threads=self._num_ffmpeg_threads, device=self._device
video_file, num_ffmpeg_threads=num_ffmpeg_threads, device=self._device
)
return decoder.get_frames_played_at(pts_list)

def get_consecutive_frames_from_video(self, video_file, numFramesToDecode):
num_ffmpeg_threads = (
int(self._num_ffmpeg_threads) if self._num_ffmpeg_threads else 0
)
decoder = VideoDecoder(
video_file, num_ffmpeg_threads=self._num_ffmpeg_threads, device=self._device
video_file, num_ffmpeg_threads=num_ffmpeg_threads, device=self._device
)
frames = []
count = 0
Expand All @@ -285,8 +320,11 @@ def get_consecutive_frames_from_video(self, video_file, numFramesToDecode):
return frames

def decode_and_transform(self, video_file, pts_list, height, width, device):
num_ffmpeg_threads = (
int(self._num_ffmpeg_threads) if self._num_ffmpeg_threads else 1
)
decoder = VideoDecoder(
video_file, num_ffmpeg_threads=self._num_ffmpeg_threads, device=self._device
video_file, num_ffmpeg_threads=num_ffmpeg_threads, device=self._device
)
frames = decoder.get_frames_played_at(pts_list)
frames = self.transforms_v2.functional.resize(frames.data, (height, width))
Expand Down Expand Up @@ -339,7 +377,9 @@ def __init__(self):

def get_frames_from_video(self, video_file, pts_list):
stream_reader = self.torchaudio.io.StreamReader(src=video_file)
stream_reader.add_basic_video_stream(frames_per_chunk=1)
stream_reader.add_basic_video_stream(
frames_per_chunk=1, decoder_option={"threads": "0"}
)
frames = []
for pts in pts_list:
stream_reader.seek(pts)
Expand All @@ -350,7 +390,9 @@ def get_frames_from_video(self, video_file, pts_list):

def get_consecutive_frames_from_video(self, video_file, numFramesToDecode):
stream_reader = self.torchaudio.io.StreamReader(src=video_file)
stream_reader.add_basic_video_stream(frames_per_chunk=1)
stream_reader.add_basic_video_stream(
frames_per_chunk=1, decoder_option={"threads": "0"}
)
frames = []
frame_cnt = 0
for vframe in stream_reader.stream():
Expand All @@ -363,7 +405,9 @@ def get_consecutive_frames_from_video(self, video_file, numFramesToDecode):

def decode_and_transform(self, video_file, pts_list, height, width, device):
stream_reader = self.torchaudio.io.StreamReader(src=video_file)
stream_reader.add_basic_video_stream(frames_per_chunk=1)
stream_reader.add_basic_video_stream(
frames_per_chunk=1, decoder_option={"threads": "1"}
)
frames = []
for pts in pts_list:
stream_reader.seek(pts)
Expand Down Expand Up @@ -634,7 +678,10 @@ def run_benchmarks(
},
label=f"video={video_file_path} {metadata_label}",
sub_label=decoder_name,
description=f"dataloader[threads={bp.num_threads},batch_size={bp.batch_size}] {num_samples} decode_and_transform()",
description=f"concurrent[threads={bp.num_threads},batch_size={bp.batch_size}] {num_samples} decode_and_transform()",
)
print(
f"{decoder_name} concurrent[threads={bp.num_threads} batch_size={bp.batch_size}]"
)
results.append(
dataloader_result.blocked_autorange(
Expand All @@ -647,7 +694,7 @@ def run_benchmarks(
decoder_name,
video_file_path,
num_samples * dataloader_parameters.batch_parameters.batch_size,
f"dataloader[threads={bp.num_threads} batch_size={bp.batch_size}] {num_samples} x decode_and_transform()",
f"concurrent[threads={bp.num_threads} batch_size={bp.batch_size}] {num_samples} x decode_and_transform()",
)
)

Expand All @@ -670,6 +717,7 @@ def run_benchmarks(
sub_label=decoder_name,
description=f"{kind} {num_samples} seek()+next()",
)
print(f"{decoder_name} {kind} {num_samples} seek()+next()")
results.append(
seeked_result.blocked_autorange(min_run_time=min_runtime_seconds)
)
Expand Down Expand Up @@ -697,6 +745,7 @@ def run_benchmarks(
sub_label=decoder_name,
description=f"batch {kind} {num_samples} seek()+next()",
)
print(f"{decoder_name} batch {kind} {num_samples} seek()+next()")
results.append(
seeked_result.blocked_autorange(
min_run_time=min_runtime_seconds
Expand Down Expand Up @@ -724,6 +773,7 @@ def run_benchmarks(
sub_label=decoder_name,
description=f"{num_consecutive_nexts} next()",
)
print(f"{decoder_name} {num_consecutive_nexts} next()")
results.append(
consecutive_frames_result.blocked_autorange(
min_run_time=min_runtime_seconds
Expand Down Expand Up @@ -753,6 +803,7 @@ def run_benchmarks(
sub_label=decoder_name,
description=f"batch {num_consecutive_nexts} next()",
)
print(f"{decoder_name} batch {num_consecutive_nexts} next()")
results.append(
consecutive_frames_result.blocked_autorange(
min_run_time=min_runtime_seconds
Expand Down
Binary file modified benchmarks/decoders/benchmark_readme_chart.png
View file
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading

0 comments on commit 1550a37

Please sign in to comment.