elm-prune 0.1.0__tar.gz

elm_prune-0.1.0/LICENSE

@@ -0,0 +1,30 @@
+MIT License
+
+Copyright (c) 2026 Iago Rodrigues
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+This license applies only to the source code, packaging files, and documentation
+contained in this repository for the Python package "mirc-dataset-handler".
+
+It does not grant any rights over the MIRC dataset itself, including but not
+limited to images, videos, annotations, derived data, redistributed copies, or
+third-party mirrors of the dataset. Any use of the MIRC dataset must follow the
+separate dataset terms, ethics requirements, and repository policies published
+in the official dataset repository and/or official distribution channels.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

elm_prune-0.1.0/PKG-INFO

@@ -0,0 +1,31 @@
+Metadata-Version: 2.4
+Name: elm-prune
+Version: 0.1.0
+Summary: ELM-based structural pruning utilities for PyTorch segmentation models.
+Author: Iago Rodrigues
+License: MIT
+Project-URL: Homepage, https://github.com/iagorichard
+Project-URL: Repository, https://github.com/iagorichard
+Project-URL: Issues, https://github.com/iagorichard
+Keywords: deep-learning,pytorch,pruning,structural-pruning,segmentation,elm
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.24
+Requires-Dist: tqdm>=4.66
+Requires-Dist: torch>=2.5
+Requires-Dist: segmentation-models-pytorch
+Requires-Dist: torch-pruning==1.6.1
+Provides-Extra: dev
+Requires-Dist: build>=1.2.2; extra == "dev"
+Requires-Dist: twine>=5.1.1; extra == "dev"
+Dynamic: license-file
+
+# elm-pruning

elm_prune-0.1.0/README.md

@@ -0,0 +1 @@
+# elm-pruning

elm_prune-0.1.0/elm_prune.egg-info/PKG-INFO

@@ -0,0 +1,31 @@
+Metadata-Version: 2.4
+Name: elm-prune
+Version: 0.1.0
+Summary: ELM-based structural pruning utilities for PyTorch segmentation models.
+Author: Iago Rodrigues
+License: MIT
+Project-URL: Homepage, https://github.com/iagorichard
+Project-URL: Repository, https://github.com/iagorichard
+Project-URL: Issues, https://github.com/iagorichard
+Keywords: deep-learning,pytorch,pruning,structural-pruning,segmentation,elm
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.24
+Requires-Dist: tqdm>=4.66
+Requires-Dist: torch>=2.5
+Requires-Dist: segmentation-models-pytorch
+Requires-Dist: torch-pruning==1.6.1
+Provides-Extra: dev
+Requires-Dist: build>=1.2.2; extra == "dev"
+Requires-Dist: twine>=5.1.1; extra == "dev"
+Dynamic: license-file
+
+# elm-pruning

elm_prune-0.1.0/elm_prune.egg-info/SOURCES.txt

@@ -0,0 +1,17 @@
+LICENSE
+README.md
+pyproject.toml
+elm_prune.egg-info/PKG-INFO
+elm_prune.egg-info/SOURCES.txt
+elm_prune.egg-info/dependency_links.txt
+elm_prune.egg-info/requires.txt
+elm_prune.egg-info/top_level.txt
+elmprune/__init__.py
+elmprune/elm.py
+elmprune/elm_importance_processor.py
+elmprune/feature_extractor.py
+elmprune/importance_processor_config.py
+elmprune/prune_config.py
+elmprune/prune_pipeline.py
+elmprune/prune_processor.py
+elmprune/utils.py

elm_prune-0.1.0/elm_prune.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

elm_prune-0.1.0/elm_prune.egg-info/requires.txt

@@ -0,0 +1,9 @@
+numpy>=1.24
+tqdm>=4.66
+torch>=2.5
+segmentation-models-pytorch
+torch-pruning==1.6.1
+
+[dev]
+build>=1.2.2
+twine>=5.1.1

elm_prune-0.1.0/elm_prune.egg-info/top_level.txt

@@ -0,0 +1 @@
+elmprune

elm_prune-0.1.0/elmprune/__init__.py

@@ -0,0 +1,20 @@
+from . import utils
+from .importance_processor_config import ImportanceProcessorConfig
+from .feature_extractor import FeatureExtractor
+from .elm_importance_processor import ELMImportanceProcessor
+from .elm import ELMRegressor
+from .prune_processor import PruneProcessor
+from .prune_config import PruneConfig
+from .prune_config import PruneVerboseLevel
+from .prune_pipeline import PrunePipeline
+
+
+__all__ = ["utils", 
+           "ImportanceProcessorConfig", 
+           "FeatureExtractor", 
+           "ELMImportanceProcessor", 
+           "ELMRegressor", 
+           "PruneProcessor", 
+           "PruneConfig",
+           "PruneVerboseLevel",
+           "PrunePipeline"]

elm_prune-0.1.0/elmprune/elm.py

@@ -0,0 +1,105 @@
+from typing import Dict, List
+import math
+import torch
+
+
+class ELMRegressor:
+
+    def __init__(self, hidden_dim: int, reg_lambda: float, activation_name: str, seed: int, eps: float, use_double_for_solver: bool):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.hidden_dim = hidden_dim
+        self.reg_lambda = reg_lambda
+        self.activation_name = activation_name
+        self.seed = seed
+        self.eps = eps
+        self.use_double_for_solver = use_double_for_solver
+
+
+    def fit(self, X: torch.Tensor, Y: torch.Tensor) -> Dict[str, torch.Tensor]:
+        X = X.to(self.device)
+        Y = Y.to(self.device)
+
+        self.X_mean = X.mean(dim=0, keepdim=True)
+        self.X_std = X.std(dim=0, keepdim=True, unbiased=False).clamp_min(self.eps)
+        Xn = (X - self.X_mean) / self.X_std
+
+        self.Y_mean = Y.mean(dim=0, keepdim=True)
+        Yn = Y - self.Y_mean
+
+        generator = torch.Generator(device=self.device)
+        generator.manual_seed(self.seed)
+
+        in_dim = Xn.shape[1]
+        out_dim = Yn.shape[1]
+
+        self.W = torch.randn((in_dim, self.hidden_dim), generator=generator, device=self.device) / math.sqrt(max(in_dim, 1))
+        self.b = torch.randn((self.hidden_dim,), generator=generator, device=self.device)
+
+        H = self.__apply_activation(Xn @ self.W + self.b, self.activation_name)
+
+        I = torch.eye(self.hidden_dim, device=self.device, dtype=H.dtype)
+        lhs = H.T @ H + self.reg_lambda * I
+        rhs = H.T @ Yn
+
+        if self.use_double_for_solver:
+            self.beta = torch.linalg.solve(lhs.double(), rhs.double()).to(H.dtype)
+        else:
+            self.beta = torch.linalg.solve(lhs, rhs)
+
+    def predict(self, X: torch.Tensor) -> torch.Tensor:
+        X = X.to(self.W.device)
+
+        Xn = (X - self.X_mean) / self.X_std
+        H = self.__apply_activation(
+            Xn @ self.W + self.b,
+            self.activation_name,
+        )
+        return H @ self.beta + self.Y_mean
+
+    def compute_ablation_importance(self, X: torch.Tensor, Y: torch.Tensor) -> List[float]:
+        """
+        Importance = increase in ELM loss when feature is neutralized to its mean value.
+        Since pruning removes less important filters, low scores should be pruned first.
+        """
+        X = X.to(self.W.device)
+        Y = Y.to(self.W.device)
+
+        base_pred = self.predict(X)
+        base_loss = self.calculate_loss(base_pred, Y).item()
+
+        importances: List[float] = []
+        X_work = X.clone()
+
+        for feature_idx in range(X.shape[1]):
+            original_column = X_work[:, feature_idx].clone()
+
+            # Neutralize feature by sending it to its mean value
+            X_work[:, feature_idx] = self.X_mean[0, feature_idx]
+
+            ablated_pred = self.predict(X_work)
+            ablated_loss = self.calculate_loss(ablated_pred, Y).item()
+
+            importance = max(ablated_loss - base_loss, 0.0)
+            importances.append(float(importance))
+
+            X_work[:, feature_idx] = original_column
+
+        return importances
+    
+    def calculate_loss(self, Y_pred, Y_original):
+        return self.__mse(Y_pred, Y_original)
+    
+    def __apply_activation(self, x: torch.Tensor, activation_name: str) -> torch.Tensor:
+        activation_name = activation_name.lower()
+
+        if activation_name == "tanh":
+            return torch.tanh(x)
+        if activation_name == "relu":
+            return torch.relu(x)
+        if activation_name == "sigmoid":
+            return torch.sigmoid(x)
+
+        raise ValueError(f"Unsupported activation: {activation_name}")
+    
+    def __mse(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        return torch.mean((pred - target) ** 2)

elm_prune-0.1.0/elmprune/elm_importance_processor.py

@@ -0,0 +1,164 @@
+from typing import Dict, Iterable, List
+from tqdm.auto import tqdm
+import torch
+from torch import nn
+from .utils import get_all_conv_layer_names, compute_constant_baseline_loss, dump_dict, load_dict
+from .elm import ELMRegressor
+from .importance_processor_config import ImportanceProcessorConfig
+from .feature_extractor import FeatureExtractor
+
+
+class ELMImportanceProcessor:
+
+    IMPORTANCES_RELATIVE_FOLDER = "importances"
+    IMPORTANCES_GLOBAL_CACHE_FILENAME = "importances_elm_global.json"
+    IMPORTANCES_LAYERWISE_CACHE_FILENAME = "importances_elm_layerwise.json"
+    IMPORTANCES_FILTERWISE_CACHE_FILENAME = "importances_elm_filterwise.json"
+
+    def __init__(self, config: ImportanceProcessorConfig, model: nn.Module, dataloader: Iterable):
+        self.config = config 
+        self.layer_names = get_all_conv_layer_names(model) if config.layer_names == "" else config.layer_names
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        self.importances_path = config.abs_path / ELMImportanceProcessor.IMPORTANCES_RELATIVE_FOLDER
+        self.importances_path.mkdir(parents=True, exist_ok=True)
+
+        self.has_global_cache, self.has_layerwise_cache, self.has_filterwise_cache = self.__verify_cache()
+        if not config.use_cache or not self.has_global_cache or not self.has_layerwise_cache or not self.has_filterwise_cache:
+            feature_extractor = FeatureExtractor(config, model, dataloader, self.layer_names)
+            self.features_by_layer, self.targets = feature_extractor.extract_feature_and_targets()
+        else:
+            print("[ELMImportanceProcessor] Not necessary to extract features here! Using cache for all importances type.")
+
+    def __verify_cache(self):
+        has_global_cache = (self.importances_path / ELMImportanceProcessor.IMPORTANCES_GLOBAL_CACHE_FILENAME).exists()
+        has_layerwise_cache = (self.importances_path / ELMImportanceProcessor.IMPORTANCES_LAYERWISE_CACHE_FILENAME).exists()
+        has_filterwise_cache = (self.importances_path / ELMImportanceProcessor.IMPORTANCES_FILTERWISE_CACHE_FILENAME).exists()
+        return has_global_cache, has_layerwise_cache, has_filterwise_cache
+
+    def compute_elm_global_importances(self) -> Dict[str, List[float]]:
+        """
+        One single ELM trained with features from all selected layers concatenated.
+        Importance of each filter = increase in ELM loss when that feature is neutralized.
+        """
+        print("[ELMImportanceProcessor] Getting importances for global...")
+
+        if self.has_global_cache:
+            print("[ELMImportanceProcessor] Not necessary to calculate importances here! Using cache for this importance type.")
+            result = load_dict(self.importances_path / ELMImportanceProcessor.IMPORTANCES_GLOBAL_CACHE_FILENAME)
+            return result
+        
+        if len(self.features_by_layer) == 0:
+            raise RuntimeError("No features were collected for ELM global importance.")
+
+        X_parts = [self.features_by_layer[layer_name] for layer_name in self.layer_names]
+        X = torch.cat(X_parts, dim=1)
+        Y = self.targets
+
+        elm_model = ELMRegressor(
+            hidden_dim=self.config.hidden_dim,
+            reg_lambda=self.config.reg_lambda,
+            activation_name=self.config.activation,
+            seed=self.config.seed,
+            eps=self.config.eps,
+            use_double_for_solver=self.config.use_double_for_solver,
+        )
+
+        elm_model.fit(X, Y)
+        importances = elm_model.compute_ablation_importance(X, Y)
+
+        result: Dict[str, List[float]] = {}
+        offset = 0
+        for layer_name in tqdm(self.layer_names, desc="ELM global feature ranking processing", dynamic_ncols=True, position=1, leave=False):
+            channels = self.features_by_layer[layer_name].shape[1]
+            result[layer_name] = importances[offset: offset + channels]
+            offset += channels
+
+        importances_dump_path = self.importances_path / ELMImportanceProcessor.IMPORTANCES_GLOBAL_CACHE_FILENAME
+        dump_dict(result, importances_dump_path)
+
+        return result
+
+    def compute_elm_layerwise_importances(self) -> Dict[str, List[float]]:
+        """
+        One ELM per layer.
+        Importance of each filter = increase in ELM loss when that feature is neutralized.
+        """
+        print("[ELMImportanceProcessor] Getting importances for layerwise...")
+
+        if self.has_layerwise_cache:
+            print("[ELMImportanceProcessor]: Not necessary to calculate importances here! Using cache for this importance type.")
+            result = load_dict(self.importances_path / ELMImportanceProcessor.IMPORTANCES_LAYERWISE_CACHE_FILENAME)
+            return result
+
+        result: Dict[str, List[float]] = {}
+        Y = self.targets.to(self.device)
+
+        for layer_name in tqdm(self.layer_names, desc="ELM layerwise feature ranking processing", dynamic_ncols=True, position=1, leave=False):
+            X = self.features_by_layer[layer_name].to(self.device)
+
+            elm_model = ELMRegressor(
+                hidden_dim=self.config.hidden_dim,
+                reg_lambda=self.config.reg_lambda,
+                activation_name=self.config.activation,
+                seed=self.config.seed,
+                eps=self.config.eps,
+                use_double_for_solver=self.config.use_double_for_solver,
+            )
+
+            elm_model.fit(X, Y)
+            importances = elm_model.compute_ablation_importance(X, Y)
+            result[layer_name] = importances
+
+        importances_dump_path = self.importances_path / ELMImportanceProcessor.IMPORTANCES_LAYERWISE_CACHE_FILENAME
+        dump_dict(result, importances_dump_path)
+
+        return result
+
+    def compute_elm_filterwise_importances(self) -> Dict[str, List[float]]:
+        """
+        One tiny ELM per filter.
+        Importance of one filter = how much that single filter alone reduces target reconstruction loss
+        compared with a constant baseline.
+        """
+        print("[ELMImportanceProcessor] Getting importances for filterwise...")
+
+        if self.has_filterwise_cache:
+            print("[ELMImportanceProcessor]: Not necessary to calculate importances here! Using cache for this importance type.")
+            result = load_dict(self.importances_path / ELMImportanceProcessor.IMPORTANCES_FILTERWISE_CACHE_FILENAME)
+            return result
+        
+        result: Dict[str, List[float]] = {}
+        Y = self.targets.to(self.device)
+        baseline_loss = compute_constant_baseline_loss(Y)
+
+        for layer_name in tqdm(self.layer_names, desc="ELM filterwise feature ranking processing", dynamic_ncols=True, position=1, leave=False):
+            X_layer = self.features_by_layer[layer_name].to(self.device)
+            layer_importances: List[float] = []
+
+            for filter_idx in range(X_layer.shape[1]):
+                X_filter = X_layer[:, filter_idx:filter_idx + 1]
+
+                elm_model = ELMRegressor(
+                    hidden_dim=self.config.hidden_dim_per_filter,
+                    reg_lambda=self.config.reg_lambda,
+                    activation_name=self.config.activation,
+                    seed=self.config.seed + filter_idx,
+                    eps=self.config.eps,
+                    use_double_for_solver=self.config.use_double_for_solver,
+                )
+
+                elm_model.fit(X_filter, Y)
+                pred = elm_model.predict(X_filter)
+                filter_loss = elm_model.calculate_loss(pred, Y).item()
+
+                # Higher reduction => more important
+                importance = max(baseline_loss - filter_loss, 0.0)
+                layer_importances.append(float(importance))
+
+            result[layer_name] = layer_importances
+
+        importances_dump_path = self.importances_path / ELMImportanceProcessor.IMPORTANCES_FILTERWISE_CACHE_FILENAME
+        dump_dict(result, importances_dump_path)
+
+        return result

elm_prune-0.1.0/elmprune/feature_extractor.py

@@ -0,0 +1,217 @@
+import sys
+import torch
+import torch.nn as nn
+from tqdm.auto import tqdm
+from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
+from .utils import get_layer_by_string
+from .importance_processor_config import ImportanceProcessorConfig
+
+
+class FeatureExtractor:
+    
+    def __init__(self, config: ImportanceProcessorConfig, model: nn.Module, dataloader: Iterable, layer_names: List[str]):
+        self.config = config 
+        self.model = model
+        self.dataloader = dataloader
+        self.layer_names = layer_names
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    def extract_feature_and_targets(self)-> Tuple[Dict[str, torch.Tensor], torch.Tensor]:
+        
+        if self.config.feature_type == "segmentation":
+            target_extractor = self.__segmentation_mask_histogram_target_extractor()
+        else:
+            target_extractor = self.__default_logits_gap_target_extractor
+
+        return self.__process_features_and_targets(target_extractor)
+    
+    def __process_features_and_targets(self, target_extractor: Callable[[Any, torch.Tensor, Optional[torch.Tensor]], torch.Tensor]) -> Tuple[Dict[str, torch.Tensor], torch.Tensor]:
+        
+        print("[FeatureExtractor] Extracting features for ELM...")
+        
+        self.model.eval()
+        self.model = self.model.to(self.device)
+        storage_device = "cpu"
+
+        hooks = []
+        feature_storage: Dict[str, List[torch.Tensor]] = {layer_name: [] for layer_name in self.layer_names}
+        target_storage: List[torch.Tensor] = []
+
+        current_batch_features: Dict[str, torch.Tensor] = {}
+
+        def make_hook(layer_name: str):
+            def hook_fn(module, inputs, output):
+                out = self.__extract_first_tensor(output)
+
+                if out.dim() == 4:
+                    # [N, C, H, W] -> [N, C]
+                    pooled = out.mean(dim=(2, 3))
+                elif out.dim() == 3:
+                    pooled = out.mean(dim=2)
+                elif out.dim() == 2:
+                    pooled = out
+                else:
+                    pooled = out.flatten(start_dim=1)
+
+                current_batch_features[layer_name] = pooled.detach().to(storage_device)
+            return hook_fn
+
+        # Register hooks
+        for layer_name in self.layer_names:
+            layer = get_layer_by_string(self.model, layer_name)
+            hooks.append(layer.register_forward_hook(make_hook(layer_name)))
+
+        try:
+            with torch.no_grad():
+                total_batches = None
+                if hasattr(self.dataloader, "__len__"):
+                    total_batches = len(self.dataloader)
+                    if self.config.max_batches is not None:
+                        total_batches = min(total_batches, self.config.max_batches)
+
+                progress_bar = tqdm(
+                    enumerate(self.dataloader),
+                    total=total_batches,
+                    desc="Collecting features",
+                    dynamic_ncols=True,
+                    file=sys.stdout,
+                    position=1, 
+                    leave=False
+                )
+
+                for batch_idx, batch in progress_bar:
+                    if self.config.max_batches is not None and batch_idx >= self.config.max_batches:
+                        break
+
+                    inputs, targets = self.__unpack_batch(batch)
+
+                    inputs = inputs.to(self.device)
+                    targets = targets.to(self.device) if targets is not None else None
+
+                    current_batch_features.clear()
+
+                    model_output = self.model(inputs)
+                    y = target_extractor(model_output, inputs, targets).to(storage_device)
+
+                    for layer_name in self.layer_names:
+                        if layer_name not in current_batch_features:
+                            raise RuntimeError(f"No activation captured for layer '{layer_name}'.")
+
+                        feature_storage[layer_name].append(current_batch_features[layer_name])
+
+                    target_storage.append(y)
+
+        finally:
+            for hook in hooks:
+                hook.remove()
+
+        features_by_layer = {
+            layer_name: torch.cat(feature_storage[layer_name], dim=0)
+            for layer_name in self.layer_names
+        }
+        targets = torch.cat(target_storage, dim=0)
+
+        return features_by_layer, targets
+
+    def __default_logits_gap_target_extractor(self, model_output: Any) -> torch.Tensor:
+        """
+        Default target extractor:
+        converts model output into [N, D] using a GAP-like reduction.
+        Very practical because it does not assume a specific task label encoding.
+        """
+        out = self.__extract_first_tensor(model_output)
+
+        if out.dim() == 4:
+            # Typical segmentation logits: [N, C, H, W] -> [N, C]
+            return out.mean(dim=(2, 3)).detach()
+        if out.dim() == 3:
+            return out.mean(dim=2).detach()
+        if out.dim() == 2:
+            return out.detach()
+
+        return out.flatten(start_dim=1).detach()
+
+    def __segmentation_mask_histogram_target_extractor(self) -> Callable[[Any, torch.Tensor, Optional[torch.Tensor]], torch.Tensor]:
+        """
+        Returns a target extractor that uses the GT mask distribution per image.
+        For each image, target becomes [p(class0), p(class1), ..., p(classK)].
+
+        This is more task-aware for segmentation.
+        """
+        def _extractor(model_output: Any, input_batch: torch.Tensor, target_batch: Optional[torch.Tensor] = None) -> torch.Tensor:
+            if target_batch is None:
+                raise ValueError("Target batch is required for mask histogram target extractor.")
+
+            y = target_batch
+
+            if y.dim() == 4 and y.shape[1] == 1:
+                y = y.squeeze(1)
+
+            if y.dim() != 3:
+                raise ValueError(f"Expected target mask with shape [N, H, W] or [N, 1, H, W], got {tuple(y.shape)}")
+
+            y = y.long()
+            histograms = []
+
+            for class_idx in range(self.config.num_classes):
+                class_ratio = (y == class_idx).float().mean(dim=(1, 2))
+                histograms.append(class_ratio)
+
+            return torch.stack(histograms, dim=1).detach()
+
+        return _extractor
+    
+    def __extract_first_tensor(self, data: Any) -> torch.Tensor:
+        if torch.is_tensor(data):
+            return data
+
+        if isinstance(data, dict):
+            if "out" in data and torch.is_tensor(data["out"]):
+                return data["out"]
+
+            for value in data.values():
+                if torch.is_tensor(value):
+                    return value
+
+        if isinstance(data, (list, tuple)):
+            for value in data:
+                if torch.is_tensor(value):
+                    return value
+
+        raise TypeError("Could not extract a torch.Tensor from model output / hook output.")
+
+    def __unpack_batch(self, batch: Any) -> Tuple[torch.torch.Tensor, Optional[torch.torch.Tensor]]:
+        if torch.is_tensor(batch):
+            return batch, None
+
+        if isinstance(batch, dict):
+            if "image" not in batch:
+                raise KeyError("Batch dict must contain key 'image'.")
+
+            x = batch["image"]
+            y = batch.get("mask", None)
+
+            if not torch.is_tensor(x):
+                raise TypeError("Batch['image'] is not a torch.Tensor.")
+
+            if y is not None and not torch.is_tensor(y):
+                y = None
+
+            return x, y
+
+        if isinstance(batch, (list, tuple)):
+            if len(batch) == 0:
+                raise ValueError("Empty batch received.")
+
+            x = batch[0]
+            y = batch[1] if len(batch) > 1 else None
+
+            if not torch.is_tensor(x):
+                raise TypeError("Batch input is not a torch.Tensor.")
+
+            if y is not None and not torch.is_tensor(y):
+                y = None
+
+            return x, y
+
+        raise TypeError(f"Unsupported batch type: {type(batch)}")

elm_prune-0.1.0/elmprune/importance_processor_config.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+from typing import Optional
+from pathlib import Path
+
+
+@dataclass
+class ImportanceProcessorConfig:
+    hidden_dim: int = 128
+    hidden_dim_per_filter: int = 16
+    reg_lambda: float = 1e-3
+    activation: str = "tanh"  # tanh | relu | sigmoid
+    max_batches: Optional[int] = None
+    eps: float = 1e-8
+    seed: int = 42
+    use_double_for_solver: bool = True
+    feature_type: str = "segmentation" # segmention | logits
+    num_classes: int = 3 # if segmentation
+    layer_names = "" # to get importance for all layers, or list (str) to specify the layer names
+    abs_path: Path = None
+    use_cache: bool = True

elm_prune-0.1.0/elmprune/prune_config.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+from typing import Literal
+from enum import Enum
+
+class PruneVerboseLevel(Enum):
+    BASIC = 1
+    BASIC_ERROR = 2
+    ALL = 3
+
+@dataclass
+class PruneConfig:
+    importance_type: Literal["elm_global", "elm_layerwise", "elm_filterwise"]
+    target_param_reduction: float          # ex.: 0.20 = reduzir 20% dos params reais
+    selection_scope: Literal["global", "local"]
+    min_channels_abs: int = 16
+    min_keep_ratio: float = 0.50
+    max_layer_prune_ratio: float = 0.35
+    per_step_layer_ratio: float = 0.05     # poda pequena por iteração
+    round_to: int = 8
+    verbose: PruneVerboseLevel = PruneVerboseLevel.BASIC_ERROR