torch-dad 0.1.0__tar.gz

torch_dad-0.1.0/PKG-INFO

@@ -0,0 +1,14 @@
+Metadata-Version: 2.4
+Name: torch-dad
+Version: 0.1.0
+Summary: A highly accelerated, backprop-free Decoupled Analytical Dense (DAD) target propagation training engine on top of PyTorch.
+Author: Mukundan Ramaswamy
+License: MIT
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0.0
+Requires-Dist: torchvision

torch_dad-0.1.0/pyproject.toml

@@ -0,0 +1,28 @@
+[build-system]
+requires = ["setuptools>=61.0.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "torch-dad"
+version = "0.1.0"
+description = "A highly accelerated, backprop-free Decoupled Analytical Dense (DAD) target propagation training engine on top of PyTorch."
+readme = "README.md"
+requires-python = ">=3.8"
+license = {text = "MIT"}
+authors = [
+    {name = "Mukundan Ramaswamy"}
+]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence"
+]
+dependencies = [
+    "torch>=2.0.0",
+    "torchvision"
+]
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["torch_dad*"]

torch_dad-0.1.0/setup.cfg

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

torch_dad-0.1.0/torch_dad/__init__.py

@@ -0,0 +1,6 @@
+from .layers import DADLinear
+from .models import DADModel
+from .trainers import DADTrainer
+
+__version__ = "0.1.0"
+__all__ = ["DADLinear", "DADModel", "DADTrainer"]

torch_dad-0.1.0/torch_dad/layers.py

@@ -0,0 +1,64 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def adamw_step_fn(p, g, m, v, lr, t):
+    """Fused tensor-operation AdamW optimization step."""
+    b1, b2, eps = 0.9, 0.999, 1e-8
+    
+    # Update moment estimates in-place
+    m.mul_(b1).add_(g, alpha=1.0 - b1)
+    v.mul_(b2).addcmul_(g, g, value=1.0 - b2)
+    
+    bias_correction1 = 1.0 - b1 ** t
+    bias_correction2 = 1.0 - b2 ** t
+    
+    step_size = lr / bias_correction1
+    denom = (v.sqrt() / torch.sqrt(bias_correction2)).add_(eps)
+    
+    # In-place weight decay and gradient descent step
+    p.mul_(1.0 - lr * 0.01)
+    p.addcdiv_(m, denom, value=-step_size)
+
+class DADLinear(nn.Module):
+    """
+    Decoupled Analytical Dense (DAD) target propagation layer.
+    Acts as a drop-in high-performance alternative to nn.Linear for backprop-free networks.
+    """
+    def __init__(self, in_features, out_features, num_classes=10, device=None):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.num_classes = num_classes
+        self.device = device
+
+        # Main Weight & Bias (No autograd gradients tracked)
+        self.W = nn.Parameter(
+            torch.randn(out_features, in_features, device=device) * (2.0 / in_features) ** 0.5,
+            requires_grad=False
+        )
+        self.bias = nn.Parameter(
+            torch.zeros(out_features, device=device),
+            requires_grad=False
+        )
+
+        # Local Task Classifier Head (Trainable)
+        self.W_loc = nn.Parameter(
+            torch.randn(num_classes, out_features, device=device) * 0.02,
+            requires_grad=False
+        )
+        self.b_loc = nn.Parameter(
+            torch.zeros(num_classes, device=device),
+            requires_grad=False
+        )
+
+        # Optimizer Moments (No Autograd tracking)
+        self.m_W = nn.Parameter(torch.zeros_like(self.W, device=device), requires_grad=False)
+        self.v_W = nn.Parameter(torch.zeros_like(self.W, device=device), requires_grad=False)
+        self.m_bias = nn.Parameter(torch.zeros_like(self.bias, device=device), requires_grad=False)
+        self.v_bias = nn.Parameter(torch.zeros_like(self.bias, device=device), requires_grad=False)
+
+    def forward(self, x):
+        z = F.linear(x, self.W, self.bias)
+        out = torch.relu(z)
+        return out, z

torch_dad-0.1.0/torch_dad/models.py

@@ -0,0 +1,33 @@
+import torch
+import torch.nn as nn
+from .layers import DADLinear
+
+class DADModel(nn.Module):
+    """
+    Base DAD Neural Network Container.
+    Inherits from nn.Module, automatically scanning and managing any nested DADLinear layers.
+    """
+    def __init__(self):
+        super().__init__()
+        
+    @property
+    def dad_layers(self):
+        """Dynamically scans and returns all DADLinear layers in the model in registration order."""
+        return [module for module in self.modules() if isinstance(module, DADLinear)]
+
+    def forward_inference(self, x):
+        """
+        Runs standard inference through the model without target propagation overhead.
+        Expects a final linear layer named `self.classifier` in subclasses.
+        """
+        x = x.view(x.size(0), -1)
+        with torch.no_grad():
+            for layer in self.dad_layers:
+                out, _ = layer(x)
+                x = out
+                
+        # Subclasses must define self.classifier (e.g. standard nn.Linear final head)
+        if hasattr(self, 'classifier'):
+            return self.classifier(x)
+        else:
+            raise AttributeError("DADModel subclasses must define self.classifier as their final output head.")

torch_dad-0.1.0/torch_dad/trainers.py

@@ -0,0 +1,120 @@
+import time
+import torch
+import torch.nn.functional as F
+from .layers import adamw_step_fn
+
+class DADTrainer:
+    """
+    Decoupled Analytical Dense (DAD) Training Engine.
+    Manages JIT step training loops, autocasting, and final closed-form classifier solving.
+    """
+    def __init__(self, model, device: torch.device):
+        self.model = model
+        self.device = device
+        
+        # Verify model has DAD layers
+        dad_layers = model.dad_layers
+        if not dad_layers:
+            raise ValueError("Provided model has no DADLinear layers registered!")
+            
+        self.num_classes = dad_layers[0].num_classes
+        self.step_counter = 0
+        self.t_tensor = torch.tensor(0.0, dtype=torch.float32, device=device)
+
+        # Dynamic JIT Compilation: JIT compile only on GPU to bypass compilation latency on CPU
+        if device.type == 'cuda':
+            self.compiled_step = torch.compile(self.unified_step)
+        else:
+            self.compiled_step = self.unified_step
+
+        # Pre-allocate closed-form solving matrices based on final layer output dimension
+        last_layer = dad_layers[-1]
+        out_features = last_layer.W.size(0)
+        self.HTH = torch.zeros(out_features + 1, out_features + 1, device=device)
+        self.HTY = torch.zeros(out_features + 1, self.num_classes, device=device)
+
+    def reset_step_counter(self):
+        """Resets the training step counters and moment updates."""
+        self.step_counter = 0
+        self.t_tensor.zero_()
+
+    def unified_step(self, x, y, lr, t):
+        """Unified, JIT-fusible forward-backward training step."""
+        # 1. Forward Pass
+        acts = [x.view(x.size(0), -1)]
+        zs = []
+        x_d = acts[0]
+        for layer in self.model.dad_layers:
+            out, z = layer(x_d)
+            acts.append(out)
+            zs.append(z)
+            x_d = out
+            
+        # 2. Decoupled Backward Pass
+        y_true = F.one_hot(y, num_classes=self.num_classes).float()
+        inv_batch_size = 1.0 / x.size(0)
+        
+        for i, layer in enumerate(self.model.dad_layers):
+            out = acts[i+1]
+            z = zs[i]
+            x_prev = acts[i]
+            
+            # Local alignment projection
+            y_pred = F.linear(out, layer.W_loc, layer.b_loc)
+            probs  = F.softmax(y_pred, dim=-1)
+            d_pred = (probs - y_true) * inv_batch_size
+
+            # Analytical local gradients calculation
+            g_W_loc = d_pred.t() @ out
+            g_b_loc = d_pred.sum(0)
+            d_h     = d_pred @ layer.W_loc
+            d_z     = d_h * (z > 0).float()
+            g_W     = d_z.t() @ x_prev
+            g_bias  = d_z.sum(0)
+
+            # In-place updates: Fused AdamW for main, ultra-fast SGD for local classifiers
+            adamw_step_fn(layer.W,     g_W,     layer.m_W,     layer.v_W,     lr, t)
+            adamw_step_fn(layer.bias,  g_bias,  layer.m_bias,  layer.v_bias,  lr, t)
+            layer.W_loc.add_(g_W_loc, alpha=-1e-3)
+            layer.b_loc.add_(g_b_loc, alpha=-1e-3)
+
+        return acts[-1]
+
+    @torch.no_grad()
+    def train_epoch(self, loader, amp_dtype, accumulate_solver=False):
+        """Trains the model for one full epoch, with optional closed-form matrix accumulations."""
+        self.model.train()
+        for x, y in loader:
+            self.step_counter += 1
+            self.t_tensor.add_(1.0)
+            
+            with torch.amp.autocast(self.device.type, dtype=amp_dtype, enabled=(self.device.type == 'cuda')):
+                out_last = self.compiled_step(x, y, 1e-3, self.t_tensor)
+                
+                # Eager solver accumulation outside JIT to guarantee zero graph breaks
+                if accumulate_solver:
+                    ones = torch.ones(out_last.size(0), 1, device=self.device)
+                    h_aug = torch.cat([out_last, ones], dim=1)
+                    y_onehot = F.one_hot(y, num_classes=self.num_classes).float()
+                    self.HTH.add_(h_aug.t() @ h_aug)
+                    self.HTY.add_(h_aug.t() @ y_onehot)
+
+    @torch.no_grad()
+    def solve_head(self, lambda_reg=1e-3):
+        """Instantly solves the optimal final classifier linear mapping in VRAM."""
+        t0 = time.time()
+        out_features = self.model.dad_layers[-1].W.size(0)
+        
+        reg = lambda_reg * torch.eye(out_features + 1, device=self.device)
+        try:
+            W_aug = torch.linalg.solve(self.HTH + reg, self.HTY)
+        except RuntimeError:
+            W_aug = torch.linalg.pinv(self.HTH + reg) @ self.HTY
+
+        W = W_aug[:-1, :].t()
+        b = W_aug[-1, :]
+        self.model.classifier.weight.copy_(W)
+        self.model.classifier.bias.copy_(b)
+        
+        elapsed = time.time() - t0
+        return elapsed

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

@@ -0,0 +1,14 @@
+Metadata-Version: 2.4
+Name: torch-dad
+Version: 0.1.0
+Summary: A highly accelerated, backprop-free Decoupled Analytical Dense (DAD) target propagation training engine on top of PyTorch.
+Author: Mukundan Ramaswamy
+License: MIT
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0.0
+Requires-Dist: torchvision

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

@@ -0,0 +1,11 @@
+README.md
+pyproject.toml
+torch_dad/__init__.py
+torch_dad/layers.py
+torch_dad/models.py
+torch_dad/trainers.py
+torch_dad.egg-info/PKG-INFO
+torch_dad.egg-info/SOURCES.txt
+torch_dad.egg-info/dependency_links.txt
+torch_dad.egg-info/requires.txt
+torch_dad.egg-info/top_level.txt

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

@@ -0,0 +1 @@
+

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

@@ -0,0 +1,2 @@
+torch>=2.0.0
+torchvision

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

@@ -0,0 +1 @@
+torch_dad