Create model_code.py

model_code.py

@@ -0,0 +1,416 @@
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+# MIXING AUGMENTATIONS
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+def alphamix_data(x, y, alpha_range=(0.3, 0.7), spatial_ratio=0.25):
+    """
+    Standard AlphaMix: Single spatially localized transparent overlay.
+    """
+    batch_size = x.size(0)
+    index = torch.randperm(batch_size, device=x.device)
+    
+    y_a, y_b = y, y[index]
+    
+    # Sample alpha from Beta distribution
+    alpha_min, alpha_max = alpha_range
+    beta_sample = torch.distributions.Beta(2.0, 2.0).sample().item()
+    alpha = alpha_min + (alpha_max - alpha_min) * beta_sample
+    
+    # Compute overlay region
+    _, _, H, W = x.shape
+    overlay_ratio = torch.sqrt(torch.tensor(spatial_ratio)).item()
+    overlay_h = int(H * overlay_ratio)
+    overlay_w = int(W * overlay_ratio)
+    
+    top = torch.randint(0, H - overlay_h + 1, (1,), device=x.device).item()
+    left = torch.randint(0, W - overlay_w + 1, (1,), device=x.device).item()
+    
+    # Blend
+    composited_x = x.clone()
+    overlay_region = alpha * x[:, :, top:top+overlay_h, left:left+overlay_w]
+    background_region = (1 - alpha) * x[index, :, top:top+overlay_h, left:left+overlay_w]
+    composited_x[:, :, top:top+overlay_h, left:left+overlay_w] = overlay_region + background_region
+    
+    return composited_x, y_a, y_b, alpha
+
+
+def alphamix_fractal(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    alpha_range=(0.3, 0.7),
+    steps_range=(1, 3),
+    triad_scales=(1/3, 1/9, 1/27),
+    beta_shape=(2.0, 2.0),
+    seed: int | None = None,
+):
+    """
+    Fractal AlphaMix: Triadic multi-patch overlays aligned to Cantor geometry.
+    Pure torch, GPU-compatible.
+    """
+    if seed is not None:
+        torch.manual_seed(seed)
+    
+    B, C, H, W = x.shape
+    device = x.device
+    
+    # Permutation for mixing
+    idx = torch.randperm(B, device=device)
+    y_a, y_b = y, y[idx]
+    
+    x_mix = x.clone()
+    total_area = H * W
+    
+    # Beta distribution for transparency sampling
+    k1, k2 = beta_shape
+    beta_dist = torch.distributions.Beta(k1, k2)
+    alpha_min, alpha_max = alpha_range
+    
+    # Storage for effective alpha calculation
+    alpha_elems = []
+    area_weights = []
+    
+    # Sample number of patches (same for all images in batch)
+    steps = torch.randint(steps_range[0], steps_range[1] + 1, (1,), device=device).item()
+    
+    for _ in range(steps):
+        # Choose triadic scale
+        scale_idx = torch.randint(0, len(triad_scales), (1,), device=device).item()
+        scale = triad_scales[scale_idx]
+        
+        # Compute patch dimensions (triadic area)
+        patch_area = max(1, int(total_area * scale))
+        side = int(torch.sqrt(torch.tensor(patch_area, dtype=torch.float32)).item())
+        h = max(1, min(H, side))
+        w = max(1, min(W, side))
+        
+        # Random position
+        top = torch.randint(0, H - h + 1, (1,), device=device).item()
+        left = torch.randint(0, W - w + 1, (1,), device=device).item()
+        
+        # Sample transparency from Beta distribution
+        alpha_raw = beta_dist.sample().item()
+        alpha = alpha_min + (alpha_max - alpha_min) * alpha_raw
+        
+        # Track for effective alpha
+        alpha_elems.append(alpha)
+        area_weights.append(h * w)
+        
+        # Blend patches
+        fg = alpha * x[:, :, top:top + h, left:left + w]
+        bg = (1 - alpha) * x[idx, :, top:top + h, left:left + w]
+        x_mix[:, :, top:top + h, left:left + w] = fg + bg
+    
+    # Compute area-weighted effective alpha
+    alpha_t = torch.tensor(alpha_elems, dtype=torch.float32, device=device)
+    area_t = torch.tensor(area_weights, dtype=torch.float32, device=device)
+    alpha_eff = (alpha_t * area_t).sum() / (area_t.sum() + 1e-12)
+    alpha_eff = alpha_eff.item()
+    
+    return x_mix, y_a, y_b, alpha_eff
+
+
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+# DEVIL'S STAIRCASE PE
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+class DevilStaircasePE(nn.Module):
+    """Devil's Staircase PE - VECTORIZED for GPU."""
+    
+    def __init__(self, levels=20, features_per_level=4, smooth_tau=0.25, base=3):
+        super().__init__()
+        self.levels = levels
+        self.features_per_level = features_per_level
+        self.tau = smooth_tau
+        self.base = base
+        
+        self.alpha = nn.Parameter(torch.tensor(0.1))
+        
+        # Precompute level scales and powers
+        self.register_buffer('k_range', torch.arange(1, levels + 1, dtype=torch.float32))
+        self.register_buffer('cantor_powers', 0.5 ** self.k_range)
+        
+        self.base_features = 2
+        if features_per_level > 2:
+            self.feature_expansion = nn.Linear(self.base_features, features_per_level)
+        else:
+            self.feature_expansion = None
+        
+    def forward(self, positions, seq_len):
+        B = positions.shape[0]
+        device = positions.device
+        
+        x = positions.float() / max(1, (seq_len - 1))
+        x = x.clamp(1e-6, 1.0 - 1e-6)  # [B]
+        
+        # VECTORIZED: Compute all levels at once
+        scales = self.base ** self.k_range.to(device)  # [levels]
+        y = (x.unsqueeze(1) * scales.unsqueeze(0)) % self.base  # [B, levels]
+        
+        # VECTORIZED: Triadic softmax for all levels
+        centers = torch.tensor([0.5, 1.5, 2.5], device=device, dtype=x.dtype)
+        d2 = (y.unsqueeze(-1) - centers) ** 2  # [B, levels, 3]
+        logits = -d2 / (self.tau + 1e-8)
+        p = F.softmax(logits, dim=-1)  # [B, levels, 3]
+        
+        # VECTORIZED: Cantor bits
+        bit_k = p[..., 2] + self.alpha * p[..., 1]  # [B, levels]
+        
+        # VECTORIZED: Cantor sum (single matmul instead of loop)
+        Cx = (bit_k * self.cantor_powers.to(device).unsqueeze(0)).sum(dim=1)  # [B]
+        
+        # VECTORIZED: Entropy and PDF
+        ent = -(p * p.clamp_min(1e-8).log()).sum(dim=-1)  # [B, levels]
+        pdf_proxy = 1.1 - ent / math.log(3.0)  # [B, levels]
+        
+        # Stack features
+        base_feat = torch.stack([bit_k, pdf_proxy], dim=-1)  # [B, levels, 2]
+        
+        if self.feature_expansion is not None:
+            # [B, levels, 2] -> [B, levels, features_per_level]
+            pe_levels = self.feature_expansion(base_feat)
+        else:
+            pe_levels = base_feat
+        
+        return pe_levels, Cx
+
+
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+# GEOMETRIC BASIN COMPATIBILITY
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+class GeometricBasinCompatibility(nn.Module):
+    """Compute geometric compatibility scores - 4-factor product."""
+    
+    def __init__(self, num_classes=100, pe_levels=20, features_per_level=4):
+        super().__init__()
+        
+        self.num_classes = num_classes
+        self.pe_levels = pe_levels
+        self.features_per_level = features_per_level
+        
+        self.class_signatures = nn.Parameter(
+            torch.randn(num_classes, pe_levels, features_per_level) * 0.1
+        )
+        
+        self.cantor_prototypes = nn.Parameter(
+            torch.linspace(0.0, 1.0, num_classes)
+        )
+        
+        self.level_resonance = nn.Parameter(
+            torch.ones(num_classes, pe_levels) / pe_levels
+        )
+    
+    def forward(self, pe_levels, cantor_measures):
+        B = pe_levels.shape[0]
+        
+        # 1. TRIADIC COMPATIBILITY
+        pe_norm = F.normalize(pe_levels, p=2, dim=-1)
+        sig_norm = F.normalize(self.class_signatures, p=2, dim=-1)
+        
+        similarities = torch.einsum('blf,clf->bcl', pe_norm, sig_norm)
+        similarities = (similarities + 1) / 2
+        
+        resonance = F.softmax(self.level_resonance, dim=-1)
+        triadic_compat = (similarities * resonance.unsqueeze(0)).sum(dim=-1)
+        
+        # 2. SELF-SIMILARITY - VECTORIZED
+        level_k = pe_levels[:, :-1, :]   # [B, 19, features] - all levels except last
+        level_k1 = pe_levels[:, 1:, :]    # [B, 19, features] - all levels except first
+
+        # Compute all pairwise similarities at once
+        sim = F.cosine_similarity(level_k, level_k1, dim=-1, eps=1e-8)  # [B, 19]
+        sim = (sim + 1) / 2
+        self_sim_pattern = sim  # No stack needed, already [B, levels-1]
+        
+        expected_patterns = torch.sigmoid(
+            self.level_resonance[:, :-1] - self.level_resonance[:, 1:]
+        )
+        
+        pattern_diff = torch.abs(
+            self_sim_pattern.unsqueeze(1) - expected_patterns.unsqueeze(0)
+        )
+        self_sim_compat = 1 - pattern_diff.mean(dim=-1)
+        self_sim_compat = torch.clamp(self_sim_compat, 0.0, 1.0)
+        
+        # 3. CANTOR COHERENCE
+        distances = torch.abs(
+            cantor_measures.unsqueeze(1) - self.cantor_prototypes.unsqueeze(0)
+        )
+        cantor_compat = torch.exp(-distances ** 2 / 0.1) + 1e-8
+        
+        # 4. HIERARCHICAL CHECK
+        split_point = self.pe_levels // 2
+        early_levels = pe_levels[:, :split_point, :].mean(dim=1)
+        late_levels = pe_levels[:, split_point:, :].mean(dim=1)
+        
+        early_targets = self.class_signatures[:, :split_point, :].mean(dim=1)
+        late_targets = self.class_signatures[:, split_point:, :].mean(dim=1)
+        
+        early_levels_norm = F.normalize(early_levels, p=2, dim=-1)
+        late_levels_norm = F.normalize(late_levels, p=2, dim=-1)
+        early_targets_norm = F.normalize(early_targets, p=2, dim=-1)
+        late_targets_norm = F.normalize(late_targets, p=2, dim=-1)
+        
+        early_compat = torch.matmul(early_levels_norm, early_targets_norm.t())
+        late_compat = torch.matmul(late_levels_norm, late_targets_norm.t())
+        
+        early_compat = (early_compat + 1) / 2
+        late_compat = (late_compat + 1) / 2
+        hier_compat = (early_compat + late_compat) / 2
+        
+        # 5. COMBINE (geometric mean)
+        eps = 1e-6
+        triadic_compat = torch.clamp(triadic_compat, eps, 1.0)
+        self_sim_compat = torch.clamp(self_sim_compat, eps, 1.0)
+        cantor_compat = torch.clamp(cantor_compat, eps, 1.0)
+        hier_compat = torch.clamp(hier_compat, eps, 1.0)
+        
+        compatibility_scores = (
+            triadic_compat * 
+            self_sim_compat * 
+            cantor_compat * 
+            hier_compat
+        ) ** 0.25
+        
+        return compatibility_scores
+
+
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+# GEOMETRIC BASIN LOSS
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+class GeometricBasinLoss(nn.Module):
+    """Loss supervising geometric basin stability field."""
+    
+    def __init__(self, temperature=0.1):
+        super().__init__()
+        self.temperature = temperature
+        
+    def forward(self, compatibility_scores, labels, mixed_labels=None, lam=None):
+        batch_size = compatibility_scores.shape[0]
+        
+        if mixed_labels is not None and lam is not None:
+            primary_compat = compatibility_scores[torch.arange(batch_size), labels]
+            secondary_compat = compatibility_scores[torch.arange(batch_size), mixed_labels]
+            
+            primary_loss = F.mse_loss(primary_compat, torch.full_like(primary_compat, lam))
+            secondary_loss = F.mse_loss(secondary_compat, torch.full_like(secondary_compat, 1 - lam))
+            
+            soft_targets = torch.zeros_like(compatibility_scores)
+            soft_targets[torch.arange(batch_size), labels] = lam
+            soft_targets[torch.arange(batch_size), mixed_labels] = 1 - lam
+            
+            compat_normalized = compatibility_scores / (compatibility_scores.sum(dim=1, keepdim=True) + 1e-8)
+            kl_loss = F.kl_div(
+                compat_normalized.log(),
+                soft_targets,
+                reduction='batchmean'
+            )
+            
+            total_loss = primary_loss + secondary_loss + 0.1 * kl_loss
+            
+        else:
+            correct_compat = compatibility_scores[torch.arange(batch_size), labels]
+            correct_loss = -torch.log(correct_compat + 1e-8).mean()
+            
+            mask = torch.ones_like(compatibility_scores)
+            mask[torch.arange(batch_size), labels] = 0
+            
+            incorrect_compat = compatibility_scores * mask
+            incorrect_loss = torch.log(1 - incorrect_compat + 1e-8).mean()
+            incorrect_loss = -incorrect_loss
+            
+            scaled_scores = compatibility_scores / self.temperature
+            log_probs = F.log_softmax(scaled_scores, dim=1)
+            contrastive_loss = F.nll_loss(log_probs, labels)
+            
+            total_loss = correct_loss + 0.5 * incorrect_loss + 0.5 * contrastive_loss
+        
+        return total_loss
+
+
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+# GEOMETRIC BASIN CLASSIFIER
+# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+class GeometricBasinClassifier(nn.Module):
+    """Geometric basin classifier with ResNet18 backbone + Cantor PE."""
+    
+    def __init__(self, num_classes=100, pe_levels=20, pe_features_per_level=4, dropout=0.1, pretrained=False):
+        super().__init__()
+        
+        self.num_classes = num_classes
+        self.pe_levels = pe_levels
+        self.pe_features_per_level = pe_features_per_level
+        
+        # ResNet18 backbone from torchvision
+        from torchvision.models import resnet18, ResNet18_Weights
+        if pretrained:
+            resnet = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
+        else:
+            resnet = resnet18(weights=None)
+        
+        # Extract feature extractor (everything except fc layer)
+        self.backbone = nn.Sequential(
+            resnet.conv1,
+            resnet.bn1,
+            resnet.relu,
+            resnet.maxpool,
+            resnet.layer1,
+            resnet.layer2,
+            resnet.layer3,
+            resnet.layer4,
+            resnet.avgpool
+        )
+        
+        # ResNet18 outputs 512 features
+        self.feature_dim = 512
+        self.dropout = nn.Dropout(dropout)
+        
+        # Devil's Staircase PE
+        self.pe = DevilStaircasePE(pe_levels, pe_features_per_level)
+        
+        # PE modulator (adjusted for ResNet18's 512 features)
+        self.pe_modulator = nn.Sequential(
+            nn.Linear(self.feature_dim, 256),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(256, pe_levels * pe_features_per_level)
+        )
+        
+        # Geometric basin
+        self.basin = GeometricBasinCompatibility(
+            num_classes, 
+            pe_levels, 
+            pe_features_per_level
+        )
+        
+    def forward(self, x, return_details=False):
+        batch_size = x.shape[0]
+        
+        # ResNet18 backbone
+        cnn_features = self.backbone(x)
+        cnn_features = torch.flatten(cnn_features, 1)
+        cnn_features = self.dropout(cnn_features)
+        
+        # Generate PE
+        positions = torch.arange(batch_size, device=x.device)
+        pe_levels, cantor_measures = self.pe(positions, seq_len=batch_size)
+        
+        # Modulate PE with CNN features
+        modulation = self.pe_modulator(cnn_features)
+        modulation = modulation.view(batch_size, self.pe_levels, self.pe_features_per_level)
+        pe_levels = pe_levels + 0.1 * modulation
+        
+        # Geometric basin compatibility
+        compatibility_scores = self.basin(pe_levels, cantor_measures)
+        
+        if return_details:
+            return {
+                'compatibility_scores': compatibility_scores,
+                'pe_levels': pe_levels,
+                'cantor_measures': cantor_measures,
+                'cnn_features': cnn_features
+            }
+        
+        return compatibility_scores