From c914029dc9258e00840a5caa4bd5832600311847 Mon Sep 17 00:00:00 2001 From: Dixing Xu Date: Tue, 3 Nov 2026 06:17:57 +0000 Subject: [PATCH] perf(entropy_models): decompress 31% faster on test_compression_ND EntropyModel.decompress converted the CDF tables with tolist() again for every string and built each output row with torch.tensor() on a long Python int list. Convert the tables once before the loop, read the decoded symbols with np.fromiter() + torch.from_numpy(), and build the indexes and dequantize over one flat spatial dim for N-D inputs. Output values, dtype and shape are unchanged. On a workload like TestEntropyBottleneck::test_compression_ND, decompress time goes from 68.6 ms to 47.4 ms. Co-Authored-By: Aiden --- compressai/entropy_models/entropy_models.py | 39 ++++++++++++++++----- 1 file changed, 30 insertions(+), 9 deletions(-) diff --git a/compressai/entropy_models/entropy_models.py b/compressai/entropy_models/entropy_models.py index 496ce5a1..84ca942a 100644 --- a/compressai/entropy_models/entropy_models.py +++ b/compressai/entropy_models/entropy_models.py @@ -27,6 +27,7 @@ # OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +import math import warnings from typing import Any, Callable, List, Optional, Tuple, Union @@ -313,17 +314,31 @@ def decompress( cdf = self._quantized_cdf outputs = cdf.new_empty(indexes.size()) + if len(strings) > 0: + # CDF tables are shared by all strings: convert them only once + cdf_list = cdf.tolist() + cdf_length = self._cdf_length.reshape(-1).int().tolist() + offset = self._offset.reshape(-1).int().tolist() + indexes = indexes.int() + for i, s in enumerate(strings): values = self.entropy_coder.decode_with_indexes( - s, - indexes[i].reshape(-1).int().tolist(), - cdf.tolist(), - self._cdf_length.reshape(-1).int().tolist(), - self._offset.reshape(-1).int().tolist(), + s, indexes[i].reshape(-1).tolist(), cdf_list, cdf_length, offset + ) + # numpy converts large int lists much faster than torch.tensor + values = np.fromiter(values, dtype=np.int32, count=len(values)) + outputs[i] = torch.from_numpy(values).reshape(outputs[i].size()) + + size = outputs.size() + if means is not None and means.dim() == len(size) > 3 and outputs.numel() > 0: + # Elementwise ops over many small dims are slow, dequantize on a + # flattened (N, C, -1) view instead + outputs = self.dequantize( + outputs.reshape(size[0], size[1], -1), + means.reshape(means.size(0), means.size(1), -1), + dtype, ) - outputs[i] = torch.tensor( - values, device=outputs.device, dtype=outputs.dtype - ).reshape(outputs[i].size()) + return outputs.reshape(size) outputs = self.dequantize(outputs, means, dtype) return outputs @@ -563,7 +578,13 @@ def compress(self, x): def decompress(self, strings, size): output_size = (len(strings), self._quantized_cdf.size(0), *size) - indexes = self._build_indexes(output_size).to(self._quantized_cdf.device) + if len(size) > 1 and all(isinstance(s, int) and s >= 0 for s in size): + # Repeating over a single flattened spatial dim is much faster + flat_size = (*output_size[:2], math.prod(size)) + indexes = self._build_indexes(flat_size).reshape(output_size) + else: + indexes = self._build_indexes(output_size) + indexes = indexes.to(self._quantized_cdf.device) medians = self._extend_ndims(self._get_medians().detach(), len(size)) medians = medians.expand(len(strings), *([-1] * (len(size) + 1))) return super().decompress(strings, indexes, medians.dtype, medians)