torch-ttt 0.0.1__tar.gz

torch_ttt-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Nikita Durasov 
+
+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.
+
+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.

torch_ttt-0.0.1/PKG-INFO

@@ -0,0 +1,16 @@
+Metadata-Version: 2.4
+Name: torch-ttt
+Version: 0.0.1
+License-File: LICENSE
+Provides-Extra: dev
+Requires-Dist: pytest-cov; extra == "dev"
+Provides-Extra: docs
+Requires-Dist: sphinxawesome-theme==6.0.0b1; extra == "docs"
+Requires-Dist: sphinx-gallery==0.18.0; extra == "docs"
+Requires-Dist: sphinxcontrib-bibtex; extra == "docs"
+Requires-Dist: sphinxcontrib-googleanalytics; extra == "docs"
+Requires-Dist: matplotlib; extra == "docs"
+Requires-Dist: tqdm; extra == "docs"
+Provides-Extra: all
+Requires-Dist: torch_ttt[dev,docs]; extra == "all"
+Dynamic: license-file

torch_ttt-0.0.1/README.md

@@ -0,0 +1,74 @@
+<div align="center">
+  <img src="docs/source/_static/images/torch-ttt.png" alt="TorchTTT" width="500">
+</div>
+
+<!-- <div style="display: flex; gap: 0px; flex-wrap: wrap; align-items: center;">
+    <a href="https://github.com/nikitadurasov/torch-ttt/stargazers" style="margin: 2px;">
+        <img src="https://img.shields.io/github/stars/nikitadurasov/torch-ttt.svg?style=social" alt="GitHub stars" style="display: inline-block; margin: 0;">
+    </a>
+    <a href="https://github.com/nikitadurasov/torch-ttt/network" style="margin: 2px;">
+        <img src="https://img.shields.io/github/forks/nikitadurasov/torch-ttt.svg?color=blue" alt="GitHub forks" style="display: inline-block; margin: 0;">
+    </a>
+    <a href="https://github.com/nikitadurasov/torch-ttt/actions/workflows/deploy-docs.yml" style="margin: 2px;">
+        <img src="https://github.com/nikitadurasov/torch-ttt/actions/workflows/deploy-docs.yml/badge.svg" alt="Documentation" style="display: inline-block; margin: 0;">
+    </a>
+</div> -->
+
+# torch-ttt
+
+**torch-ttt** is a package designed to work with [Test-Time Training (TTT)](https://arxiv.org/abs/1909.13231) techniques and make your networks more generalizable. It aims to be modular, easy to integrate into existing pipelines, and collaborative— including as many methods as possible. Reach out to add yours!
+
+<p align="center">
+    >> You can find our webpage and documentation here:</strong> 
+    <a href="https://torch-ttt.github.io">torch-ttt.github.io</a>
+</p>
+
+> **torch-ttt** is in its early stages, so changes are expected. Contributions are welcome—feel free to get involved! If you run into any bugs or issues, don’t hesitate to submit an issue.
+
+---
+
+This package provides a streamlined API for a variety of TTT methods through *Engines*, which are lightweight wrappers around your model. These Engines are:
+
+- **Easy to use** – The internal logic of TTT methods is fully encapsulated, so you only need to wrap your model with an Engine, and you're ready to go.  
+- **Highly modular and standardized** – Each Engine follows the same interface, allowing methods to be used interchangeably, making it easy to find the best fit for your application.  
+- **Minimal changes required** – Enabling a TTT method for your model requires only a few additional lines of code.  
+
+Check out the [Quick Start](https://torch-ttt.github.io/quickstart.html) guide or the [API reference](https://torch-ttt.github.io/api.html) for a more detailed explanation of how Engines work and their core concepts.
+
+# Installation
+
+**torch-ttt** requires Python 3.10 or greater. Install the desired PyTorch version in your environment. 
+
+For the latest development version you can run,
+
+```console
+pip install git+https://github.com/nikitadurasov/torch-ttt.git
+```
+
+While we do not support PyPI yet, support is expected very soon!
+
+# Quickstart
+
+We provide a **Quick Start** guide to help you get started smoothly. Take a look here: [Quick Start](https://torch-ttt.github.io/quickstart.html) and see how to integrate TTT methods into your project.
+
+<!-- # Implemented TTTs
+
+## Baselines
+
+
+| TTT-Method                                     | Image | Text | Graph | Audio |
+|-----------------------------------------------|:----------:|:--------------:|:------------:|:---------------------:|
+| [TTT](https://arxiv.org/abs/1909.13231)                      |     ⏳     |       ⏳       |      ⏳      |          ⏳           |
+| [MaskedTTT](https://arxiv.org/abs/2209.07522)                      |     ⏳     |       ⏳       |      ⏳      |          ⏳           |
+| [TTT++](https://proceedings.neurips.cc/paper/2021/hash/b618c3210e934362ac261db280128c22-Abstract.html)                      |     ⏳     |       ⏳       |      ⏳      |          ⏳           |
+| [ActMAD](https://arxiv.org/abs/2211.12870)                      |     ⏳     |       ⏳       |      ⏳      |          ⏳           |
+| [SHOT](https://arxiv.org/abs/2002.08546)                      |     ⏳     |       ⏳       |      ⏳      |          ⏳           |
+| [TENT](https://arxiv.org/abs/2006.10726)                      |     ⏳     |       ⏳       |      ⏳      |          ⏳           | -->
+
+# Tutorials
+
+We offer a variety of tutorials to help users gain a deeper understanding of the implemented methods and see how they can be applied to different tasks. Visit the [Tutorials](https://torch-ttt.github.io/auto_examples/index.html) page to explore them.
+
+# Documentation
+
+Our aim is to provide comprehensive documentation for all included TTT methods, covering their theoretical foundations, practical benefits, and efficient integration into your project. We also offer tutorials that illustrate their applications and guide you through their effective usage.

torch_ttt-0.0.1/pyproject.toml

@@ -0,0 +1,49 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "torch-ttt"
+version = "0.0.1"
+
+[project.optional-dependencies]
+
+dev = [
+    "pytest-cov",
+]
+
+docs = [
+    "sphinxawesome-theme==6.0.0b1",
+    "sphinx-gallery==0.18.0",
+    "sphinxcontrib-bibtex",
+    "sphinxcontrib-googleanalytics",
+    "matplotlib",
+    "tqdm"
+]
+
+all = [
+    "torch_ttt[dev, docs]"
+]
+
+[tool.ruff]
+line-length = 100
+target-version = "py310"
+
+exclude = [
+    "build",         # Ignore the 'build' directory
+    "dist",          # Ignore the 'dist' directory
+    "__pycache__",   # Ignore Python's cache directories
+    "notebooks"
+]
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["torch_ttt*"]
+
+[tool.ruff.lint.pydocstyle]
+convention = "google"
+
+[tool.coverage.run]
+branch = true
+source = ["torch_ttt"]
+omit = ["*/tests/*"]

torch_ttt-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

torch_ttt-0.0.1/tests/test_imports.py

@@ -0,0 +1,32 @@
+# ruff: noqa: F401
+import unittest
+
+
+class TestImports(unittest.TestCase):
+    def test_import_loss_registry(self):
+        try:
+            from torch_ttt.loss_registry import LossRegistry
+        except ImportError as e:
+            self.fail(f"Failed to import LossRegistry: {e}")
+
+    def test_import_engine_registry(self):
+        try:
+            from torch_ttt.engine_registry import EngineRegistry
+        except ImportError as e:
+            self.fail(f"Failed to import EngineRegistry: {e}")
+
+    def test_import_base_loss(self):
+        try:
+            from torch_ttt.loss.base_loss import BaseLoss
+        except ImportError as e:
+            self.fail(f"Failed to import BaseLoss: {e}")
+
+    def test_import_base_engine(self):
+        try:
+            from torch_ttt.engine.base_engine import BaseEngine
+        except ImportError as e:
+            self.fail(f"Failed to import BaseEngine: {e}")
+
+
+if __name__ == "__main__":
+    unittest.main()

torch_ttt-0.0.1/torch_ttt/engine/__init__.py

@@ -0,0 +1,5 @@
+# ruff: noqa: F401
+# from .base_engine import BaseEngine
+# from .ttt_engine import TTTEngine
+# from .ttt_pp_engine import TTTPPEngine
+# from .masked_ttt_engine import MaskedTTTEngine

torch_ttt-0.0.1/torch_ttt/engine/actmad_engine.py

@@ -0,0 +1,174 @@
+import torch
+from typing import List, Dict, Any, Tuple, Union
+from contextlib import contextmanager
+
+from torch.utils.data import DataLoader
+from torch_ttt.engine.base_engine import BaseEngine
+from torch_ttt.engine_registry import EngineRegistry
+
+__all__ = ["ActMADEngine"]
+
+
+# TODO: add cuda support
+@EngineRegistry.register("actmad")
+class ActMADEngine(BaseEngine):
+    """**ActMAD** approach: multi-level pixel-wise feature alignment.
+    
+    Args:
+        model (torch.nn.Module): Model to be trained with TTT.
+        features_layer_names (List[str] | str): List of layer names to be used for feature alignment.
+        optimization_parameters (dict): The optimization parameters for the engine.
+
+    :Example:
+
+    .. code-block:: python
+
+        from torch_ttt.engine.actmad_engine import ActMADEngine
+
+        model = MyModel()
+        engine = ActMADEngine(model, ["fc1", "fc2"])
+        optimizer = torch.optim.Adam(engine.parameters(), lr=1e-4)
+
+        # Training
+        engine.train()
+        for inputs, labels in train_loader:
+            optimizer.zero_grad()
+            outputs, loss_ttt = engine(inputs)
+            loss = criterion(outputs, labels) + alpha * loss_ttt
+            loss.backward()
+            optimizer.step()
+
+        # Compute statistics for features alignment
+        engine.compute_statistics(train_loader)
+
+        # Inference
+        engine.eval()
+        for inputs, labels in test_loader:
+            output, loss_ttt = engine(inputs)
+
+    Reference:
+
+        "ActMAD: Activation Matching to Align Distributions for Test-Time Training", M. Jehanzeb Mirza, Pol Jane Soneira, Wei Lin, Mateusz Kozinski, Horst Possegger, Horst Bischof
+
+        Paper link: `PDF <https://proceedings.neurips.cc/paper/2021/hash/b618c3210e934362ac261db280128c22-Abstract.html>`_
+    """
+
+    def __init__(
+            self,
+            model: torch.nn.Module,
+            features_layer_names: Union[List[str], str],
+            optimization_parameters: Dict[str, Any] = {},
+    ):
+        super().__init__()
+        self.model = model
+        self.features_layer_names = features_layer_names
+        self.optimization_parameters = optimization_parameters
+
+        if isinstance(features_layer_names, str):
+            self.features_layer_names = [features_layer_names]
+
+        # TODO: rewrite this
+        self.target_modules = []
+        for layer_name in self.features_layer_names:
+            layer_exists = False
+            for name, module in model.named_modules():
+                if name == layer_name:
+                    layer_exists = True
+                    self.target_modules.append(module)
+                    break
+            if not layer_exists:
+                raise ValueError(f"Layer {layer_name} does not exist in the model.")
+            
+        self.reference_mean = None
+        self.reference_var = None
+
+    def ttt_forward(self, inputs) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass of the model.
+
+        Args:
+            inputs (torch.Tensor): Input tensor.
+
+        Returns:
+            Returns the current model prediction and rotation loss value.
+        """
+        with self.__capture_hook() as features_hooks:
+            outputs = self.model(inputs)
+            features = [hook.output for hook in features_hooks]
+
+        # don't compute loss during training
+        if self.training:
+            return outputs, 0
+        
+        if self.reference_var is None or self.reference_mean is None:
+            raise ValueError(
+                "Reference statistics are not computed. Please call `compute_statistics` method."
+            )
+
+        l1_loss = torch.nn.L1Loss(reduction='mean')
+        features_means = [torch.mean(feature, dim=0) for feature in features]
+        features_vars = [torch.var(feature, dim=0) for feature in features]
+
+        loss = 0
+        for i in range(len(self.target_modules)):
+            print(features_means[i].device, self.reference_mean[i].device)
+            loss += l1_loss(features_means[i], self.reference_mean[i])
+            loss += l1_loss(features_vars[i], self.reference_var[i])
+
+        return outputs, loss
+    
+    def compute_statistics(self, dataloader: DataLoader) -> None:
+        """Extract and compute reference statistics for features.
+
+        Args:
+            dataloader (DataLoader): The dataloader used for extracting features. It can return tuples of tensors, with the first element expected to be the input tensor.
+
+        Raises:
+            ValueError: If the dataloader is empty or features have mismatched dimensions.
+        """
+        
+        self.model.eval()
+        feat_stack = [[] for _ in self.target_modules]
+
+        # TODO: compute variance in more memory efficient way
+        with torch.no_grad():
+            device = next(self.model.parameters()).device
+            for sample in dataloader:
+                if len(sample) < 1:
+                    raise ValueError("Dataloader returned an empty batch.")
+                
+                inputs = sample[0].to(device)
+                with self.__capture_hook() as features_hooks:
+                    _ = self.model(inputs)
+                    features = [hook.output.cpu() for hook in features_hooks]
+
+                for i, feature in enumerate(features):
+                    feat_stack[i].append(feature)
+
+        # Compute mean and variance 
+        self.reference_mean = [torch.mean(torch.cat(feat), dim=0).to(device) for feat in feat_stack]
+        self.reference_var = [torch.var(torch.cat(feat), dim=0).to(device) for feat in feat_stack]
+
+    @contextmanager
+    def __capture_hook(self):
+        """Context manager to capture features via a forward hook."""
+    
+        class OutputHook:
+            def __init__(self):
+                self.output = None
+
+            def hook(self, module, input, output):
+                self.output = output
+
+        hook_handels = []
+        features_hooks = []
+        for module in self.target_modules:
+            hook = OutputHook()
+            features_hooks.append(hook)
+            hook_handle = module.register_forward_hook(hook.hook)
+            hook_handels.append(hook_handle)
+
+        try:
+            yield features_hooks
+        finally:
+            for hook in hook_handels:
+                hook.remove()

torch_ttt-0.0.1/torch_ttt/engine/base_engine.py

@@ -0,0 +1,47 @@
+from abc import ABC, abstractmethod
+from typing import Tuple
+import torch.nn as nn
+import torch
+from copy import deepcopy
+
+class BaseEngine(nn.Module, ABC):
+
+    def __init__(self):
+        nn.Module.__init__(self)
+        self.optimization_parameters = {}
+
+    @abstractmethod
+    def ttt_forward(self, inputs) -> Tuple[torch.Tensor, torch.Tensor]:
+        pass
+
+    def forward(self, inputs):
+
+        if self.training:
+            return self.ttt_forward(inputs)
+
+        # TODO: optimization pipeline should be more flexible and 
+        # user-defined, need some special structure for that
+        optimization_parameters = self.optimization_parameters or {} 
+
+        optimizer_name = optimization_parameters.get("optimizer_name", "adam")
+        num_steps = optimization_parameters.get("num_steps", 3)
+        lr = optimization_parameters.get("lr", 1e-2)
+        copy_model = optimization_parameters.get("copy_model", "True")
+
+        running_engine = deepcopy(self) if copy_model else self
+
+        parameters = filter(lambda p: p.requires_grad, running_engine.model.parameters())
+        if optimizer_name == "adam":
+            optimizer = torch.optim.Adam(parameters, lr=lr)
+
+        loss = 0 # default value
+        for _ in range(num_steps):
+            optimizer.zero_grad()
+            _, loss = running_engine.ttt_forward(inputs)
+            loss.backward()
+            optimizer.step()
+
+        with torch.no_grad():
+            final_outputs = running_engine.model(inputs)
+
+        return final_outputs, loss

torch_ttt-0.0.1/torch_ttt/engine/masked_ttt_engine.py

@@ -0,0 +1,16 @@
+from torch_ttt.engine.base_engine import BaseEngine
+from torch_ttt.engine_registry import EngineRegistry
+
+__all__ = ["MaskedTTTEngine"]
+
+
+# TODO: add cuda support
+@EngineRegistry.register("masked_ttt")
+class MaskedTTTEngine(BaseEngine):
+    """Masked autoencoders-based **test-time training** approach."""
+
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, inputs):
+        pass

torch_ttt-0.0.1/torch_ttt/engine/tent_engine.py

@@ -0,0 +1,49 @@
+import torch
+from typing import Dict, Any, Tuple
+from torch_ttt.engine.base_engine import BaseEngine
+from torch_ttt.engine_registry import EngineRegistry
+
+__all__ = ["TentEngine"]
+
+@EngineRegistry.register("tent")
+class TentEngine(BaseEngine):
+    """**TENT**: Fully test-time adaptation by entropy minimization.
+
+    Args:
+        model (torch.nn.Module): The model to adapt.
+        optimization_parameters (dict): Hyperparameters for adaptation.
+
+    Reference:
+        "TENT: Fully Test-Time Adaptation by Entropy Minimization"
+        Dequan Wang, Evan Shelhamer, et al.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        optimization_parameters: Dict[str, Any] = {},
+    ):
+        super().__init__()
+        self.model = model
+        self.optimization_parameters = optimization_parameters
+
+        # Tent adapts only affine parameters in BatchNorm
+        self.model.train()
+        self._configure_bn()
+
+    def _configure_bn(self):
+        for module in self.model.modules():
+            if isinstance(module, torch.nn.BatchNorm2d):
+                module.requires_grad_(True)
+                module.track_running_stats = False
+            else:
+                for param in module.parameters(recurse=False):
+                    param.requires_grad = False
+
+    def ttt_forward(self, inputs) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass and entropy loss computation."""
+        outputs = self.model(inputs)
+        probs = torch.nn.functional.softmax(outputs, dim=1)
+        log_probs = torch.nn.functional.log_softmax(outputs, dim=1)
+        entropy = -torch.sum(probs * log_probs, dim=1).mean()
+        return outputs, entropy