From 1e8030788ec42481cada1c2d6db4f56d836b2f8d Mon Sep 17 00:00:00 2001 From: Heeji Han Date: Fri, 24 Jul 2026 21:36:28 +0000 Subject: [PATCH 1/3] [feat] update total runtime as the sum of module runtimes modified: compressai_vision/pipelines/split_inference/image_split_inference.py modified: compressai_vision/pipelines/split_inference/video_split_inference.py --- .../pipelines/split_inference/image_split_inference.py | 10 ++++++++-- .../pipelines/split_inference/video_split_inference.py | 10 ++++++++-- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/compressai_vision/pipelines/split_inference/image_split_inference.py b/compressai_vision/pipelines/split_inference/image_split_inference.py index a2b99d1..beb7edc 100644 --- a/compressai_vision/pipelines/split_inference/image_split_inference.py +++ b/compressai_vision/pipelines/split_inference/image_split_inference.py @@ -142,7 +142,10 @@ def __call__( self.bitstream_name, file_prefix, ) - self.update_time_elapsed("encode", (time_measure() - start)) + elapsed_enc_time = time_measure() - start + total_enc_runtime = sum(enc_time_by_module.values()) + total_enc_time = total_enc_runtime if total_enc_runtime is not 0 else elapsed_enc_time + self.update_time_elapsed("encode", (total_enc_time)) if self.is_mac_calculation: self.acc_kmac_and_pixels_info( "feature_reduction", enc_complexity[0], enc_complexity[1] @@ -183,7 +186,10 @@ def __call__( dec_features, dec_time_by_module, dec_complexity = self._decompress( codec, res["bitstream"], self.codec_output_dir, file_prefix ) - self.update_time_elapsed("decode", (time_measure() - start)) + elapsed_dec_time = time_measure() - start + total_dec_runtime = sum(dec_time_by_module.values()) + total_dec_time = total_dec_runtime if total_dec_runtime is not 0 else elapsed_dec_time + self.update_time_elapsed("decode", (total_dec_time)) if self.is_mac_calculation: self.acc_kmac_and_pixels_info( "feature_restoration", dec_complexity[0], dec_complexity[1] diff --git a/compressai_vision/pipelines/split_inference/video_split_inference.py b/compressai_vision/pipelines/split_inference/video_split_inference.py index 6734011..e3b42a9 100644 --- a/compressai_vision/pipelines/split_inference/video_split_inference.py +++ b/compressai_vision/pipelines/split_inference/video_split_inference.py @@ -247,7 +247,10 @@ def __call__( res, enc_time_by_module, enc_complexity = self._compress( codec, features, self.codec_output_dir, self.bitstream_name, "" ) - self.update_time_elapsed("encode", (time_measure() - start)) + elapsed_enc_time = time_measure() - start + total_enc_runtime = sum(enc_time_by_module.values()) + total_enc_time = total_enc_runtime if total_enc_runtime is not 0 else elapsed_enc_time + self.update_time_elapsed("encode", (total_enc_time)) self.add_time_details("encode", enc_time_by_module) if self.is_mac_calculation: self.add_kmac_and_pixels_info( @@ -286,7 +289,10 @@ def __call__( dec_features, dec_time_by_module, dec_complexity = self._decompress( codec, res["bitstream"], self.codec_output_dir, "" ) - self.update_time_elapsed("decode", (time_measure() - start)) + elapsed_dec_time = time_measure() - start + total_dec_runtime = sum(dec_time_by_module.values()) + total_dec_time = total_dec_runtime if total_dec_runtime is not 0 else elapsed_dec_time + self.update_time_elapsed("decode", (total_dec_time)) self.add_time_details("decode", dec_time_by_module) if self.is_mac_calculation: self.add_kmac_and_pixels_info( From bb13a06173b1d350d538af485fe6c8a67a00fc3a Mon Sep 17 00:00:00 2001 From: Heeji Han Date: Fri, 24 Jul 2026 22:18:50 +0000 Subject: [PATCH 2/3] [fix] fix memory leak during OIV6 MAC calculation --- compressai_vision/utils/measure_complexity.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/compressai_vision/utils/measure_complexity.py b/compressai_vision/utils/measure_complexity.py index d6a7f33..8798623 100644 --- a/compressai_vision/utils/measure_complexity.py +++ b/compressai_vision/utils/measure_complexity.py @@ -203,6 +203,24 @@ def _flops_to_kmacs(total_flops: float) -> float: return float(total_flops) / 1e3 +# (mem-fix) KMACs are deterministic w.r.t. (module, input shapes), so trace each +# shape only once. Per-image FlopCountAnalysis (torch.jit tracing) retains feature-map +# tensors and steadily grows CPU RAM (OOM); caching avoids that and speeds up +# repeated measurement. +_KMACS_CACHE: dict = {} + + +def _shape_key(x): + """Build a hashable shape key from inputs (possibly nested).""" + if torch.is_tensor(x): + return ("t", tuple(x.shape)) + if isinstance(x, (list, tuple)): + return ("s", tuple(_shape_key(v) for v in x)) + if isinstance(x, dict): + return ("d", tuple((k, _shape_key(v)) for k, v in x.items())) + return ("o", repr(x)) + + def measure_kmacs(module: nn.Module, inputs, tag: str = None) -> float: """ Measure KMACs for a given module using fvcore. @@ -249,6 +267,15 @@ def _cast(x): elif not isinstance(inputs, tuple): inputs = (inputs,) # safe fallback + # (mem-fix) Shape-based cache lookup: skip re-tracing for the same (module, input shapes). + cache_key = None + try: + cache_key = (module.__class__.__qualname__, id(p), _shape_key(inputs)) + except Exception: + cache_key = None + if cache_key is not None and cache_key in _KMACS_CACHE: + return _KMACS_CACHE[cache_key] + with torch.no_grad(): flops = FlopCountAnalysis(module, inputs) flops.set_op_handle( @@ -280,6 +307,8 @@ def _cast(x): kmacs = _flops_to_kmacs(total_flops) name = tag or module.__class__.__name__ # print(f"[INFO] {name}: KMACs = {kmacs}") + if cache_key is not None: + _KMACS_CACHE[cache_key] = kmacs return kmacs From d7a9208e4615d3378be68011da1f1f8b66a0b65f Mon Sep 17 00:00:00 2001 From: Fabien Racape Date: Mon, 27 Jul 2026 23:56:07 +0200 Subject: [PATCH 3/3] fix: linting --- .../pipelines/split_inference/image_split_inference.py | 8 ++++++-- .../pipelines/split_inference/video_split_inference.py | 8 ++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/compressai_vision/pipelines/split_inference/image_split_inference.py b/compressai_vision/pipelines/split_inference/image_split_inference.py index beb7edc..13b4a3c 100644 --- a/compressai_vision/pipelines/split_inference/image_split_inference.py +++ b/compressai_vision/pipelines/split_inference/image_split_inference.py @@ -144,7 +144,9 @@ def __call__( ) elapsed_enc_time = time_measure() - start total_enc_runtime = sum(enc_time_by_module.values()) - total_enc_time = total_enc_runtime if total_enc_runtime is not 0 else elapsed_enc_time + total_enc_time = ( + total_enc_runtime if total_enc_runtime != 0 else elapsed_enc_time + ) self.update_time_elapsed("encode", (total_enc_time)) if self.is_mac_calculation: self.acc_kmac_and_pixels_info( @@ -188,7 +190,9 @@ def __call__( ) elapsed_dec_time = time_measure() - start total_dec_runtime = sum(dec_time_by_module.values()) - total_dec_time = total_dec_runtime if total_dec_runtime is not 0 else elapsed_dec_time + total_dec_time = ( + total_dec_runtime if total_dec_runtime != 0 else elapsed_dec_time + ) self.update_time_elapsed("decode", (total_dec_time)) if self.is_mac_calculation: self.acc_kmac_and_pixels_info( diff --git a/compressai_vision/pipelines/split_inference/video_split_inference.py b/compressai_vision/pipelines/split_inference/video_split_inference.py index e3b42a9..f6f8d1a 100644 --- a/compressai_vision/pipelines/split_inference/video_split_inference.py +++ b/compressai_vision/pipelines/split_inference/video_split_inference.py @@ -249,7 +249,9 @@ def __call__( ) elapsed_enc_time = time_measure() - start total_enc_runtime = sum(enc_time_by_module.values()) - total_enc_time = total_enc_runtime if total_enc_runtime is not 0 else elapsed_enc_time + total_enc_time = ( + total_enc_runtime if total_enc_runtime != 0 else elapsed_enc_time + ) self.update_time_elapsed("encode", (total_enc_time)) self.add_time_details("encode", enc_time_by_module) if self.is_mac_calculation: @@ -291,7 +293,9 @@ def __call__( ) elapsed_dec_time = time_measure() - start total_dec_runtime = sum(dec_time_by_module.values()) - total_dec_time = total_dec_runtime if total_dec_runtime is not 0 else elapsed_dec_time + total_dec_time = ( + total_dec_runtime if total_dec_runtime != 0 else elapsed_dec_time + ) self.update_time_elapsed("decode", (total_dec_time)) self.add_time_details("decode", dec_time_by_module) if self.is_mac_calculation: