Create make_chart_1.py

make_chart_1.py

@@ -0,0 +1,420 @@
+#@title MobiusNet Aggregate Analysis - 1024 Samples
+!pip install -q datasets safetensors huggingface_hub
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from datasets import load_dataset
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file as load_safetensors
+import matplotlib.pyplot as plt
+import numpy as np
+from sklearn.decomposition import PCA
+from typing import Tuple
+from collections import defaultdict
+import math
+import json
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# ============================================================================
+# MOBIUSNET (compact)
+# ============================================================================
+
+class MobiusLens(nn.Module):
+    def __init__(self, dim, layer_idx, total_layers, scale_range=(1.0, 9.0)):
+        super().__init__()
+        self.t = layer_idx / max(total_layers - 1, 1)
+        scale_span = scale_range[1] - scale_range[0]
+        step = scale_span / max(total_layers, 1)
+        self.register_buffer('scales', torch.tensor([scale_range[0] + self.t * scale_span,
+                                                      scale_range[0] + self.t * scale_span + step]))
+        self.twist_in_angle = nn.Parameter(torch.tensor(self.t * math.pi))
+        self.twist_in_proj = nn.Linear(dim, dim, bias=False)
+        self.omega = nn.Parameter(torch.tensor(math.pi))
+        self.alpha = nn.Parameter(torch.tensor(1.5))
+        self.phase_l, self.drift_l = nn.Parameter(torch.zeros(2)), nn.Parameter(torch.ones(2))
+        self.phase_m, self.drift_m = nn.Parameter(torch.zeros(2)), nn.Parameter(torch.zeros(2))
+        self.phase_r, self.drift_r = nn.Parameter(torch.zeros(2)), nn.Parameter(-torch.ones(2))
+        self.accum_weights = nn.Parameter(torch.tensor([0.4, 0.2, 0.4]))
+        self.xor_weight = nn.Parameter(torch.tensor(0.7))
+        self.gate_norm = nn.LayerNorm(dim)
+        self.twist_out_angle = nn.Parameter(torch.tensor(-self.t * math.pi))
+        self.twist_out_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x):
+        cos_t, sin_t = torch.cos(self.twist_in_angle), torch.sin(self.twist_in_angle)
+        x = x * cos_t + self.twist_in_proj(x) * sin_t
+        x_norm = torch.tanh(x)
+        t = x_norm.abs().mean(dim=-1, keepdim=True).unsqueeze(-2)
+        x_exp = x_norm.unsqueeze(-2)
+        s = self.scales.view(-1, 1)
+        def wave(phase, drift):
+            a = self.alpha.abs() + 0.1
+            pos = s * self.omega * (x_exp + drift.view(-1, 1) * t) + phase.view(-1, 1)
+            return torch.exp(-a * torch.sin(pos).pow(2)).prod(dim=-2)
+        L, M, R = wave(self.phase_l, self.drift_l), wave(self.phase_m, self.drift_m), wave(self.phase_r, self.drift_r)
+        w = torch.softmax(self.accum_weights, dim=0)
+        xor_w = torch.sigmoid(self.xor_weight)
+        lr = xor_w * (L + R - 2*L*R).abs() + (1 - xor_w) * L * R
+        gate = torch.sigmoid(self.gate_norm((w[0]*L + w[1]*M + w[2]*R) * (0.5 + 0.5*lr)))
+        x = x * gate
+        cos_t, sin_t = torch.cos(self.twist_out_angle), torch.sin(self.twist_out_angle)
+        return x * cos_t + self.twist_out_proj(x) * sin_t, gate
+
+class MobiusConvBlock(nn.Module):
+    def __init__(self, channels, layer_idx, total_layers, scale_range=(1.0, 9.0), reduction=0.5):
+        super().__init__()
+        self.conv = nn.Sequential(nn.Conv2d(channels, channels, 3, padding=1, groups=channels, bias=False),
+                                   nn.Conv2d(channels, channels, 1, bias=False), nn.BatchNorm2d(channels))
+        self.lens = MobiusLens(channels, layer_idx, total_layers, scale_range)
+        third = channels // 3
+        which_third = layer_idx % 3
+        mask = torch.ones(channels)
+        mask[which_third*third : which_third*third + third + (channels%3 if which_third==2 else 0)] = reduction
+        self.register_buffer('thirds_mask', mask.view(1, -1, 1, 1))
+        self.residual_weight = nn.Parameter(torch.tensor(0.9))
+
+    def forward(self, x):
+        identity = x
+        h = self.conv(x).permute(0, 2, 3, 1)
+        h, gate = self.lens(h)
+        h = h.permute(0, 3, 1, 2) * self.thirds_mask
+        rw = torch.sigmoid(self.residual_weight)
+        return rw * identity + (1 - rw) * h, gate
+
+class MobiusNet(nn.Module):
+    def __init__(self, in_chans=1, num_classes=1000, channels=(64,128,256),
+                 depths=(2,2,2), scale_range=(0.5,2.5), use_integrator=True):
+        super().__init__()
+        total_layers = sum(depths)
+        channels = list(channels)
+        self.stem = nn.Sequential(nn.Conv2d(in_chans, channels[0], 3, padding=1, bias=False), nn.BatchNorm2d(channels[0]))
+        self.stages = nn.ModuleList()
+        self.downsamples = nn.ModuleList()
+        layer_idx = 0
+        for si, d in enumerate(depths):
+            self.stages.append(nn.ModuleList([MobiusConvBlock(channels[si], layer_idx+i, total_layers, scale_range) for i in range(d)]))
+            layer_idx += d
+            if si < len(depths)-1:
+                self.downsamples.append(nn.Sequential(nn.Conv2d(channels[si], channels[si+1], 3, stride=2, padding=1, bias=False), nn.BatchNorm2d(channels[si+1])))
+        self.integrator = nn.Sequential(nn.Conv2d(channels[-1], channels[-1], 3, padding=1, bias=False), nn.BatchNorm2d(channels[-1]), nn.GELU()) if use_integrator else nn.Identity()
+        self.pool = nn.AdaptiveAvgPool2d(1)
+        self.head = nn.Linear(channels[-1], num_classes)
+
+    def forward_with_intermediates(self, x):
+        out = {'input': x, 'stem': None, 'stages': [], 'gates': [], 'final': None}
+        x = self.stem(x)
+        out['stem'] = x
+        for i, stage in enumerate(self.stages):
+            acts, gates = [], []
+            for block in stage:
+                x, g = block(x)
+                acts.append(x)
+                gates.append(g)
+            out['stages'].append(acts)
+            out['gates'].append(gates)
+            if i < len(self.downsamples):
+                x = self.downsamples[i](x)
+        x = self.integrator(x)
+        out['final'] = x
+        return self.head(self.pool(x).flatten(1)), out
+
+# ============================================================================
+# LOAD MODEL
+# ============================================================================
+
+print("Loading model...")
+config_path = hf_hub_download("AbstractPhil/mobiusnet-distillations",
+    "checkpoints/mobius_tiny_s_imagenet_clip_vit_l14/20260111_000512/config.json")
+with open(config_path) as f:
+    config = json.load(f)
+model_path = hf_hub_download("AbstractPhil/mobiusnet-distillations",
+    "checkpoints/mobius_tiny_s_imagenet_clip_vit_l14/20260111_000512/checkpoints/best_model.safetensors")
+
+cfg = config['model']
+model = MobiusNet(cfg['in_chans'], cfg['num_classes'], tuple(cfg['channels']),
+                  tuple(cfg['depths']), tuple(cfg['scale_range']), cfg['use_integrator']).to(device)
+model.load_state_dict(load_safetensors(model_path))
+model.eval()
+print(f"✓ Loaded MobiusNet")
+
+# ============================================================================
+# AGGREGATE OVER 1024 SAMPLES
+# ============================================================================
+
+print("\nProcessing 1024 samples...")
+ds = load_dataset("AbstractPhil/imagenet-clip-features-orderly", "clip_vit_l14",
+                  split="validation", streaming=True).with_format("torch")
+loader = DataLoader(ds, batch_size=64)
+
+# Accumulators
+n_samples = 0
+total_correct = 0
+agg = {
+    'input': {'sum': None, 'sum_sq': None},
+    'stem': {'sum': None, 'sum_sq': None},
+    'final': {'sum': None, 'sum_sq': None},
+}
+gate_stats = defaultdict(lambda: {'sum': 0, 'sum_sq': 0, 'min': float('inf'), 'max': float('-inf'), 'count': 0})
+stage_stats = defaultdict(lambda: {'sum': None, 'sum_sq': None})
+
+# Class-wise gate means
+class_gate_means = defaultdict(lambda: defaultdict(list))
+
+for batch_idx, batch in enumerate(loader):
+    if n_samples >= 1024:
+        break
+    
+    features = batch['clip_features'].view(-1, 1, 24, 32).to(device)
+    labels = batch['label'].to(device)
+    bs = features.shape[0]
+    
+    with torch.no_grad():
+        logits, intermediates = model.forward_with_intermediates(features)
+        preds = logits.argmax(dim=-1)
+        total_correct += (preds == labels).sum().item()
+    
+    # Aggregate inputs, stem, final
+    for key in ['input', 'stem', 'final']:
+        tensor = intermediates[key].detach()
+        if agg[key]['sum'] is None:
+            agg[key]['sum'] = tensor.sum(dim=0)
+            agg[key]['sum_sq'] = (tensor ** 2).sum(dim=0)
+        else:
+            agg[key]['sum'] += tensor.sum(dim=0)
+            agg[key]['sum_sq'] += (tensor ** 2).sum(dim=0)
+    
+    # Aggregate gates and stages
+    for si, (acts, gates) in enumerate(zip(intermediates['stages'], intermediates['gates'])):
+        for bi, (act, gate) in enumerate(zip(acts, gates)):
+            key = f"S{si}B{bi}"
+            
+            # Gate stats
+            g = gate.detach()
+            gate_stats[key]['sum'] += g.mean().item() * bs
+            gate_stats[key]['sum_sq'] += (g.mean(dim=(1,2,3)) ** 2).sum().item()
+            gate_stats[key]['min'] = min(gate_stats[key]['min'], g.min().item())
+            gate_stats[key]['max'] = max(gate_stats[key]['max'], g.max().item())
+            gate_stats[key]['count'] += bs
+            
+            # Stage activation stats
+            a = act.detach()
+            if stage_stats[key]['sum'] is None:
+                stage_stats[key]['sum'] = a.sum(dim=0)
+                stage_stats[key]['sum_sq'] = (a ** 2).sum(dim=0)
+            else:
+                stage_stats[key]['sum'] += a.sum(dim=0)
+                stage_stats[key]['sum_sq'] += (a ** 2).sum(dim=0)
+            
+            # Per-class gate means
+            for i in range(bs):
+                lbl = labels[i].item()
+                class_gate_means[key][lbl].append(g[i].mean().item())
+    
+    n_samples += bs
+    if (batch_idx + 1) % 4 == 0:
+        print(f"  Processed {n_samples} samples...")
+
+print(f"\n✓ Processed {n_samples} samples")
+print(f"✓ Overall accuracy: {total_correct / n_samples:.2%}")
+
+# ============================================================================
+# COMPUTE MEANS AND STDS
+# ============================================================================
+
+def compute_mean_std(agg_dict, n):
+    mean = agg_dict['sum'] / n
+    std = torch.sqrt(agg_dict['sum_sq'] / n - mean ** 2 + 1e-8)
+    return mean.cpu().numpy(), std.cpu().numpy()
+
+input_mean, input_std = compute_mean_std(agg['input'], n_samples)
+stem_mean, stem_std = compute_mean_std(agg['stem'], n_samples)
+final_mean, final_std = compute_mean_std(agg['final'], n_samples)
+
+stage_means, stage_stds = {}, {}
+for key in stage_stats:
+    stage_means[key], stage_stds[key] = compute_mean_std(stage_stats[key], n_samples)
+
+# ============================================================================
+# VISUALIZATION
+# ============================================================================
+
+fig = plt.figure(figsize=(24, 18))
+
+# Row 1: Input/Stem aggregates
+ax = fig.add_subplot(4, 6, 1)
+ax.imshow(input_mean[0], cmap='viridis', aspect='auto')
+ax.set_title(f"Input Mean\n[1×24×32]", fontsize=10)
+ax.axis('off')
+
+ax = fig.add_subplot(4, 6, 2)
+ax.imshow(input_std[0], cmap='magma', aspect='auto')
+ax.set_title(f"Input Std\nσ={input_std.mean():.4f}", fontsize=10)
+ax.axis('off')
+
+ax = fig.add_subplot(4, 6, 3)
+pca = PCA(n_components=3)
+stem_pca = pca.fit_transform(stem_mean.reshape(64, -1).T).reshape(24, 32, 3)
+stem_pca = (stem_pca - stem_pca.min()) / (stem_pca.max() - stem_pca.min() + 1e-8)
+ax.imshow(stem_pca, aspect='auto')
+ax.set_title(f"Stem Mean PCA\n({pca.explained_variance_ratio_.sum()*100:.1f}%)", fontsize=10)
+ax.axis('off')
+
+ax = fig.add_subplot(4, 6, 4)
+ax.imshow(stem_std.mean(axis=0), cmap='magma', aspect='auto')
+ax.set_title(f"Stem Std (avg ch)\nσ={stem_std.mean():.4f}", fontsize=10)
+ax.axis('off')
+
+ax = fig.add_subplot(4, 6, 5)
+pca = PCA(n_components=3)
+final_pca = pca.fit_transform(final_mean.reshape(256, -1).T).reshape(6, 8, 3)
+final_pca = (final_pca - final_pca.min()) / (final_pca.max() - final_pca.min() + 1e-8)
+ax.imshow(final_pca, aspect='auto')
+ax.set_title(f"Final Mean PCA\n({pca.explained_variance_ratio_.sum()*100:.1f}%)", fontsize=10)
+ax.axis('off')
+
+ax = fig.add_subplot(4, 6, 6)
+ax.imshow(np.linalg.norm(final_mean, axis=0), cmap='hot', aspect='auto')
+ax.set_title(f"Final Mean L2\nμ={np.linalg.norm(final_mean, axis=0).mean():.2f}", fontsize=10)
+ax.axis('off')
+
+# Row 2: Stage means
+for i, key in enumerate(['S0B0', 'S0B1', 'S1B0', 'S1B1', 'S2B0', 'S2B1']):
+    ax = fig.add_subplot(4, 6, 7 + i)
+    m = stage_means[key]
+    pca = PCA(n_components=3)
+    m_pca = pca.fit_transform(m.reshape(m.shape[0], -1).T).reshape(m.shape[1], m.shape[2], 3)
+    m_pca = (m_pca - m_pca.min()) / (m_pca.max() - m_pca.min() + 1e-8)
+    ax.imshow(m_pca, aspect='auto')
+    ax.set_title(f"{key} Mean\n[{m.shape[0]}×{m.shape[1]}×{m.shape[2]}]", fontsize=10)
+    ax.axis('off')
+
+# Row 3: Stage stds + Gate summary
+for i, key in enumerate(['S0B0', 'S0B1', 'S1B0', 'S1B1', 'S2B0', 'S2B1']):
+    ax = fig.add_subplot(4, 6, 13 + i)
+    s = stage_stds[key]
+    ax.imshow(s.mean(axis=0), cmap='magma', aspect='auto')
+    gs = gate_stats[key]
+    gate_mean = gs['sum'] / gs['count']
+    ax.set_title(f"{key} Std | Gate μ={gate_mean:.3f}\nrange=[{gs['min']:.2f}, {gs['max']:.2f}]", fontsize=9)
+    ax.axis('off')
+
+# Row 4: Analysis plots
+ax = fig.add_subplot(4, 6, 19)
+keys = ['S0B0', 'S0B1', 'S1B0', 'S1B1', 'S2B0', 'S2B1']
+gate_means_plot = [gate_stats[k]['sum'] / gate_stats[k]['count'] for k in keys]
+gate_mins = [gate_stats[k]['min'] for k in keys]
+gate_maxs = [gate_stats[k]['max'] for k in keys]
+x = np.arange(len(keys))
+ax.bar(x, gate_means_plot, color='steelblue', alpha=0.7)
+ax.errorbar(x, gate_means_plot, yerr=[np.array(gate_means_plot) - gate_mins, np.array(gate_maxs) - gate_means_plot],
+            fmt='none', color='black', capsize=3)
+ax.set_xticks(x)
+ax.set_xticklabels(keys, fontsize=9)
+ax.set_ylabel('Gate Activation')
+ax.set_title('Gate Means (± range)', fontsize=10)
+ax.set_ylim(0, 1)
+
+# Lens parameters
+ax = fig.add_subplot(4, 6, 20)
+lens_data = []
+for si, stage in enumerate(model.stages):
+    for bi, block in enumerate(stage):
+        L = block.lens
+        lens_data.append({
+            'name': f"S{si}B{bi}",
+            'omega': L.omega.item(),
+            'alpha': L.alpha.item(),
+            'xor': torch.sigmoid(L.xor_weight).item()
+        })
+omegas = [d['omega'] for d in lens_data]
+alphas = [d['alpha'] for d in lens_data]
+xors = [d['xor'] for d in lens_data]
+x = np.arange(len(lens_data))
+width = 0.25
+ax.bar(x - width, omegas, width, label='ω', alpha=0.7)
+ax.bar(x, alphas, width, label='α', alpha=0.7)
+ax.bar(x + width, xors, width, label='xor', alpha=0.7)
+ax.set_xticks(x)
+ax.set_xticklabels([d['name'] for d in lens_data], fontsize=9)
+ax.legend(fontsize=8)
+ax.set_title('Lens Parameters', fontsize=10)
+
+# Distribution flow
+ax = fig.add_subplot(4, 6, 21)
+flow_labels = ['Input', 'Stem'] + keys + ['Final']
+flow_stds = [input_std.std(), stem_std.std()] + [stage_stds[k].std() for k in keys] + [final_std.std()]
+ax.plot(flow_labels, flow_stds, 'o-', color='purple', linewidth=2, markersize=8)
+ax.set_ylabel('Std of Std')
+ax.set_title('Activation Variance Flow', fontsize=10)
+ax.tick_params(axis='x', rotation=45)
+
+# Per-class gate variance (are gates class-specific?)
+ax = fig.add_subplot(4, 6, 22)
+class_variance = []
+for key in keys:
+    per_class = class_gate_means[key]
+    class_means = [np.mean(v) for v in per_class.values() if len(v) > 0]
+    class_variance.append(np.std(class_means) if len(class_means) > 1 else 0)
+ax.bar(keys, class_variance, color='coral', alpha=0.7)
+ax.set_ylabel('Std of Class Gate Means')
+ax.set_title('Gate Class-Specificity', fontsize=10)
+ax.tick_params(axis='x', rotation=45)
+
+# Accuracy bar
+ax = fig.add_subplot(4, 6, 23)
+ax.bar(['Accuracy'], [total_correct / n_samples], color='green', alpha=0.7)
+ax.set_ylim(0, 1)
+ax.set_title(f'Overall: {total_correct / n_samples:.1%}\n({total_correct}/{n_samples})', fontsize=10)
+
+# Summary text
+ax = fig.add_subplot(4, 6, 24)
+summary = f"""AGGREGATE STATS (n={n_samples})
+
+Input:  μ={input_mean.mean():.6f}, σ={input_std.mean():.6f}
+Stem:   μ={stem_mean.mean():.6f}, σ={stem_std.mean():.4f}
+Final:  μ={final_mean.mean():.4f}, σ={final_std.mean():.4f}
+
+Gate Progression:
+  S0: {gate_means_plot[0]:.3f} → {gate_means_plot[1]:.3f}
+  S1: {gate_means_plot[2]:.3f} → {gate_means_plot[3]:.3f}
+  S2: {gate_means_plot[4]:.3f} → {gate_means_plot[5]:.3f}
+
+Accuracy: {total_correct / n_samples:.2%}
+"""
+ax.text(0.05, 0.95, summary, transform=ax.transAxes, fontsize=9,
+        va='top', family='monospace')
+ax.set_title('Summary', fontsize=10)
+ax.axis('off')
+
+plt.suptitle(f'MobiusNet Aggregate Analysis: {n_samples} Samples from CLIP-ViT-L14 Features', fontsize=14, fontweight='bold')
+plt.tight_layout()
+plt.savefig("mobiusnet_aggregate_1024.png", dpi=150, bbox_inches="tight")
+plt.show()
+
+# ============================================================================
+# PRINT DETAILED STATS
+# ============================================================================
+
+print(f"\n{'='*70}")
+print(f"AGGREGATE STATISTICS ({n_samples} samples)")
+print(f"{'='*70}")
+
+print(f"\nActivation Flow:")
+print(f"  Input:  μ={input_mean.mean():.6f}, σ={input_std.mean():.6f}")
+print(f"  Stem:   μ={stem_mean.mean():.6f}, σ={stem_std.mean():.6f}")
+for key in keys:
+    m, s = stage_means[key], stage_stds[key]
+    print(f"  {key}:   μ={m.mean():.6f}, σ={s.mean():.6f}")
+print(f"  Final:  μ={final_mean.mean():.6f}, σ={final_std.mean():.6f}")
+
+print(f"\nGate Statistics:")
+for key in keys:
+    gs = gate_stats[key]
+    print(f"  {key}: μ={gs['sum']/gs['count']:.4f}, range=[{gs['min']:.3f}, {gs['max']:.3f}]")
+
+print(f"\nClass-Specificity (std of per-class gate means):")
+for i, key in enumerate(keys):
+    print(f"  {key}: {class_variance[i]:.6f}")