activationscope 0.1.0__tar.gz

activationscope-0.1.0/MANIFEST.in

@@ -0,0 +1,5 @@
+include README.md
+include pyproject.toml
+include setup.py
+recursive-include csrc *.cpp *.hpp
+recursive-include activationscope *.py *.pyi

activationscope-0.1.0/PKG-INFO

@@ -0,0 +1,94 @@
+Metadata-Version: 2.4
+Name: activationscope
+Version: 0.1.0
+Summary: High-performance PyTorch plugin for tracking and analyzing intermediate neural network activations
+Author: Jan Miksa
+Project-URL: Homepage, https://github.com/OneAndZero24/ActivationScope
+Project-URL: Repository, https://github.com/OneAndZero24/ActivationScope
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: C++
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0.1
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Requires-Dist: pytest-benchmark; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: build; extra == "dev"
+
+# ActivationScope
+
+***Jan Miksa @ IDEAS Research Institute***
+
+**High-performance PyTorch activation tracker with online reduction functionality for efficient model analysis.**
+
+Built on Python + C++ with native `libtorch` hooks and **TorchScript** (`torch.jit.script`) reductions compiled to `.pt` files.
+
+**Key Benefits**
+- Zero‑copy read‑back: activation tensors are shared between C++ and Python without extra copies.
+- Native C++ hooks: no Python compute overhead per forward pass.
+- Flexible policy knobs (`StoragePolicy`, `ReductionPolicy`, `CapturePolicy`, `CaptureMode`) let you balance memory, compute, and I/O.
+- Direct-to-disk streaming (`StoragePolicy.DISK`) — activations are written directly from C++ to disk. Ideal for long-running training loops with very large models. Activations are read back on demand from `.dat` files.
+- Works with large models (e.g., diffusion) and supports streaming statistics for online use cases.
+
+---
+
+## Quick Start
+
+Every tracked layer stores full activations by default — no registration needed:
+
+```python
+import activationscope
+
+with activationscope.ActivationScope().track(model) as tracker:
+    for x, y in dataloader:
+        out = model(x)
+        loss.backward()
+
+acts = tracker.activations  # {layer_name: [Tensor, ...]} across all batches
+```
+
+## Performance
+
+### Toy model — 48 × Linear(256,256), batch=32, 200 forwards, CPU
+
+| Approach | ms/forward | Overhead vs baseline | Data captured |
+|---|---|---|---|
+| No tracking | 2.05 | — | — |
+| Naive Python hooks | 3.13 | +52.7% | 594 MiB |
+| **ActivationScope** | **2.65** | **+29.2%** | **594 MiB** |
+
+- **Peak VMS identical** — Scope 402,506 vs Naive 402,630 MiB (~0.03% diff, within ASLR noise)
+- **1.18× faster** than naive Python hooks (3.13 → 2.65 ms/fwd)
+- **95 layers tracked** (inputs + outputs across 48 linear layers)
+- **Zero-copy readback**: 594 MiB in 2.4 ms
+
+Run it yourself:
+```bash
+# Toy model (fast, GPU or CPU)
+PYTHONPATH=. python -m benchmark.runner
+
+# Pretrained ResNet-18 (requires torchvision)
+PYTHONPATH=. python -m benchmark.runner --model resnet18
+```
+
+---
+
+## Usage Guide
+
+For detailed usage instructions, see the [Usage Documentation](docs/usage/README.md).
+
+## Development Guide
+
+Documentation and developer setup information is available in [Development Documentation](docs/development/README.md).
+
+## Design Documentation
+
+The design document outlining the architecture and implementation details can be found in [Design Documentation](docs/DESIGN.md).

activationscope-0.1.0/README.md

@@ -0,0 +1,69 @@
+# ActivationScope
+
+***Jan Miksa @ IDEAS Research Institute***
+
+**High-performance PyTorch activation tracker with online reduction functionality for efficient model analysis.**
+
+Built on Python + C++ with native `libtorch` hooks and **TorchScript** (`torch.jit.script`) reductions compiled to `.pt` files.
+
+**Key Benefits**
+- Zero‑copy read‑back: activation tensors are shared between C++ and Python without extra copies.
+- Native C++ hooks: no Python compute overhead per forward pass.
+- Flexible policy knobs (`StoragePolicy`, `ReductionPolicy`, `CapturePolicy`, `CaptureMode`) let you balance memory, compute, and I/O.
+- Direct-to-disk streaming (`StoragePolicy.DISK`) — activations are written directly from C++ to disk. Ideal for long-running training loops with very large models. Activations are read back on demand from `.dat` files.
+- Works with large models (e.g., diffusion) and supports streaming statistics for online use cases.
+
+---
+
+## Quick Start
+
+Every tracked layer stores full activations by default — no registration needed:
+
+```python
+import activationscope
+
+with activationscope.ActivationScope().track(model) as tracker:
+    for x, y in dataloader:
+        out = model(x)
+        loss.backward()
+
+acts = tracker.activations  # {layer_name: [Tensor, ...]} across all batches
+```
+
+## Performance
+
+### Toy model — 48 × Linear(256,256), batch=32, 200 forwards, CPU
+
+| Approach | ms/forward | Overhead vs baseline | Data captured |
+|---|---|---|---|
+| No tracking | 2.05 | — | — |
+| Naive Python hooks | 3.13 | +52.7% | 594 MiB |
+| **ActivationScope** | **2.65** | **+29.2%** | **594 MiB** |
+
+- **Peak VMS identical** — Scope 402,506 vs Naive 402,630 MiB (~0.03% diff, within ASLR noise)
+- **1.18× faster** than naive Python hooks (3.13 → 2.65 ms/fwd)
+- **95 layers tracked** (inputs + outputs across 48 linear layers)
+- **Zero-copy readback**: 594 MiB in 2.4 ms
+
+Run it yourself:
+```bash
+# Toy model (fast, GPU or CPU)
+PYTHONPATH=. python -m benchmark.runner
+
+# Pretrained ResNet-18 (requires torchvision)
+PYTHONPATH=. python -m benchmark.runner --model resnet18
+```
+
+---
+
+## Usage Guide
+
+For detailed usage instructions, see the [Usage Documentation](docs/usage/README.md).
+
+## Development Guide
+
+Documentation and developer setup information is available in [Development Documentation](docs/development/README.md).
+
+## Design Documentation
+
+The design document outlining the architecture and implementation details can be found in [Design Documentation](docs/DESIGN.md).

activationscope-0.1.0/activationscope/_C.pyi

@@ -0,0 +1,46 @@
+"""Type stubs for the compiled C++ extension ``activationscope._C``.
+
+TorchScript reduction path — reductions are compiled via torch.jit.script,
+serialised to .pt files, and loaded by the C++ backend.
+"""
+
+from typing import Any, Dict, List
+
+import torch
+
+# ── Session lifecycle ──────────────────────────────────────────────
+
+def session_create(
+    storage: int,
+    reduction: int,
+    sample_every: int,
+    max_batches: int,
+    auto_cpu_threshold_bytes: int,
+    use_pinned: bool,
+    session_dir: str = "",
+    capture_mode: int = 0,
+) -> int: ...
+
+def session_destroy(id: int) -> None: ...
+
+def session_readback(
+    id: int,
+) -> Dict[str, List[torch.Tensor]]: ...
+
+def session_readback_disk(
+    id: int,
+) -> Dict[str, List[str]]: ...
+
+def session_clear(id: int) -> None: ...
+
+def session_detach_hooks(id: int) -> None: ...
+
+# ── Hook registration ──────────────────────────────────────────────
+
+def session_register_hooks(
+    id: int,
+    module_ptr: Any,
+    layer_key: str,
+    capture_dir_int: int,
+    reduction_path: str = "",
+) -> None: ...

activationscope-0.1.0/activationscope/__init__.py

@@ -0,0 +1,12 @@
+"""Public API for ActivationScope — session-scoped, zero-copy, native hooks."""
+
+from activationscope.policies import StoragePolicy, ReductionPolicy, CapturePolicy, CaptureMode
+from activationscope.tracker import ActivationScope
+
+__all__ = [
+    "ActivationScope",
+    "StoragePolicy",
+    "ReductionPolicy",
+    "CapturePolicy",
+    "CaptureMode",
+]

activationscope-0.1.0/activationscope/_naive.py

@@ -0,0 +1,105 @@
+"""Shared naive PyTorch forward-hook tracker for benchmarking and testing.
+
+Provides a simple reference implementation of activation accumulation using
+Python ``register_forward_hook``.  Used by the benchmark runner to compare
+memory/throughput against ActivationScope's C++ hooks, and by parity tests
+to verify equivalent behavior.
+"""
+
+import contextlib
+from fnmatch import fnmatch
+
+import torch
+
+from activationscope.policies import CaptureMode
+
+
+class NaiveHookTracker:
+    """Accumulate per-layer activations via Python forward hooks.
+
+    Mirrors the behavior of ``ActivationScope`` with ``STORE_ALL`` reduction
+    and ``EVERY`` capture policy, but uses pure-Python hooks.  Not intended
+    for production use — exists solely as a reference for benchmarking and
+    correctness testing.
+
+    Parameters
+    ----------
+    capture_mode : CaptureMode
+        Controls copy behaviour of captured tensors:
+        * ``CaptureMode.REFERENCE`` (default) — ``.detach().cpu()``.
+          Tensors share storage with the autograd graph and may be
+          invalidated on subsequent forward passes.
+        * ``CaptureMode.SNAPSHOT`` — ``.detach().cpu().clone()``.
+          Completely independent copies, safe to keep across forwards.
+
+    Usage::
+
+        n = NaiveHookTracker(capture_mode=CaptureMode.SNAPSHOT)
+        with n.track(model):
+            for _ in range(n_batches):
+                _ = model(x)
+        acts = n.activations  # dict[str, list[Tensor]]
+    """
+
+    def __init__(self, capture_mode: CaptureMode = CaptureMode.REFERENCE):
+        self._handles: list = []
+        self._activations: dict[str, list[torch.Tensor]] = {}
+        self._capture_mode = capture_mode
+
+    def track(self, model, layers=None):
+        """Context manager: attach → yield self → detach."""
+
+        @contextlib.contextmanager
+        def _ctx():
+            self._attach(model, layers=layers)
+            yield self
+            self._detach()
+
+        return _ctx()
+
+    def _attach(self, model, layers=None):
+        containers = (torch.nn.ModuleList, torch.nn.ModuleDict, torch.nn.Sequential)
+        all_modules = {
+            n: m
+            for n, m in model.named_modules()
+            if n != "" and not isinstance(m, containers)
+        }
+        targets = (
+            all_modules
+            if layers is None
+            else {
+                n: m
+                for n, m in all_modules.items()
+                if any(fnmatch(n, p) for p in layers)
+            }
+        )
+        self._activations = {name: [] for name in targets}
+        for name in targets:
+            handle = targets[name].register_forward_hook(self._make_hook(name))
+            self._handles.append(handle)
+
+    def _make_hook(self, layer_name):
+        capture_mode = self._capture_mode
+
+        def hook_fn(_module, _inp, out):
+            if isinstance(out, torch.Tensor):
+                t = out.detach().cpu()
+                self._activations[layer_name].append(
+                    t.clone() if capture_mode == CaptureMode.SNAPSHOT else t
+                )
+            elif isinstance(out, (tuple, list)) and len(out) > 0:
+                t = out[0].detach().cpu()
+                self._activations[layer_name].append(
+                    t.clone() if capture_mode == CaptureMode.SNAPSHOT else t
+                )
+
+        return hook_fn
+
+    def _detach(self):
+        for h in self._handles:
+            h.remove()
+        self._handles = []
+
+    @property
+    def activations(self):
+        return self._activations

activationscope-0.1.0/activationscope/policies.py

@@ -0,0 +1,58 @@
+"""ActivationScope policy enumerations.
+
+Three independently tunable policy knobs that govern every aspect of memory
+and compute behaviour.  Integer-valued to pass trivially to the C++ backend.
+"""
+
+
+class _PolicyMeta(type):
+    """Metaclass that makes policy classes iterable and subscriptable."""
+
+    def __iter__(cls):
+        return iter(cls._members_)
+
+    def __class_getitem__(cls, item):
+        return cls._members_[item]
+
+
+class StoragePolicy(int, metaclass=_PolicyMeta):
+    """Where tensor data lives after capture."""
+
+    AUTO = 0    # Heuristic: < threshold → CPU, >= threshold → GPU
+    CPU  = 1    # Transfer to host memory
+    GPU  = 2    # Stay on original device
+    DISK = 3    # Stream directly to disk; bypass in-memory accumulation
+
+    _members_ = (AUTO, CPU, GPU, DISK)
+
+
+class ReductionPolicy(int):
+    """What gets kept vs. reduced across batches."""
+
+    STORE_ALL  = 0  # Full tensor per batch appended
+    STREAMING  = 1  # Per-batch reduction output replaces/accumulates in-place
+    FINAL_ONLY = 2  # Last-batch activation overwrites previous
+
+
+class CapturePolicy(int, metaclass=_PolicyMeta):
+    """When and how often hooks fire."""
+
+    EVERY    = 0  # Every forward fires hooks
+    SAMPLE_N = 1  # Captures every Nth forward
+    MAX_K    = 2  # Captures exactly K batches then stops
+
+    _members_ = (EVERY, SAMPLE_N, MAX_K)
+
+
+class CaptureMode(int, metaclass=_PolicyMeta):
+    """Whether to clone captured tensors for independent storage.
+
+    Controls the copy behaviour of both the native C++ tracker
+    (``activationscope.tracker.ActivationScope``) and the pure‑Python
+    fallback (``activationscope._naive.NaiveHookTracker``).
+    """
+
+    REFERENCE = 0  # detach() only — shares storage with autograd graph
+    SNAPSHOT  = 1  # detach() + clone() — completely independent copy
+
+    _members_ = (REFERENCE, SNAPSHOT)