[fix] preprocessing features prior to nn part2 before dumping

Author: jianwensong

compressai_vision/pipelines/base.py

@@ -32,7 +32,6 @@
 import logging
 import os
 
-from collections import OrderedDict
 from enum import Enum
 from pathlib import Path
 from typing import Callable, Dict, Tuple
@@ -50,7 +49,11 @@
     min_max_normalization,
 )
 from compressai_vision.model_wrappers import BaseWrapper
-from compressai_vision.utils import FileLikeHasher, freeze_zip_timestamps
+from compressai_vision.utils import (
+    FileLikeHasher,
+    contiguous_features,
+    freeze_zip_timestamps,
+)
 
 
 class Parts(Enum):
@@ -404,11 +407,7 @@ def _from_features_to_output(
             self.logger.debug(f"dumping features prior to nn part2 in: {feature_dir}")
 
             # [TODO] align with nn_task_part1 dump features
-            features_to_dump = {
-                k: v
-                for k, v in sorted(x.items(), key=lambda kv: str(kv[0]))
-                if not k.startswith("file")
-            }
+            features_to_dump = contiguous_features(x)
 
             with freeze_zip_timestamps():
                 if dump_feature_hash:

compressai_vision/utils/__init__.py

@@ -29,7 +29,7 @@
 
 from . import dataio, git, pip, system
 from .external_exec import get_max_num_cpus
-from .hash import FileLikeHasher, freeze_zip_timestamps
+from .hash import FileLikeHasher, contiguous_features, freeze_zip_timestamps
 from .misc import dict_sum, dl_to_ld, ld_to_dl, metric_tracking, time_measure, to_cpu
 
 __all__ = [
@@ -46,4 +46,5 @@
     "ld_to_dl",
     "FileLikeHasher",
     "freeze_zip_timestamps",
+    "contiguous_features",
 ]

compressai_vision/utils/hash.py

@@ -31,9 +31,13 @@
 import hashlib
 import zipfile
 
+from collections import OrderedDict
+from collections.abc import Mapping, Sequence
 from contextlib import contextmanager
 from typing import Tuple
 
+import torch
+
 
 class FileLikeHasher:
     def __init__(self, fn, algo: str = "md5"):
@@ -70,3 +74,23 @@ def _patched(self, *args, **kwargs):
         yield
     finally:
         zipfile.ZipInfo.__init__ = _orig_init
+
+
+def contiguous_features(obj):
+    if isinstance(obj, torch.Tensor):
+        return obj.to("cpu").contiguous().clone()
+
+    if isinstance(obj, Mapping):
+        return OrderedDict(
+            (k, contiguous_features(v))
+            for k, v in sorted(obj.items(), key=lambda item: str(item[0]))
+            if not str(k).startswith("file")
+        )
+
+    if isinstance(obj, set):
+        return tuple(sorted(obj, key=str))
+
+    if isinstance(obj, Sequence) and not isinstance(obj, (str, bytes)):
+        return type(obj)(contiguous_features(v) for v in obj)
+
+    return obj