Add files using upload-large-folder tool

test/nn/MoE4Own-batch-00034-20260331-223020-b2fdc7ac534972535c85e467f80ad728.py

@@ -0,0 +1,503 @@
+# Auto-generated 4-Expert Heterogeneous MoE: AirNet + AlexNet + DPN107 + SqueezeNet-2
+# Four-expert mixture with learned gating, mixup, and LR scheduling
+
+import numpy as np
+import torch.nn.functional as F
+from torchvision.models._api import WeightsEnum
+from torchvision.models._utils import _ovewrite_named_param
+from typing import Any, Optional
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+import torch.optim as optim
+
+
+def supported_hyperparameters():
+    return {'dropout', 'lr', 'momentum'}
+
+
+
+# ============================================================================
+# EXPERT 1: AirNet
+# ============================================================================
+
+
+class AirInitBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.layers = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class AirUnit(nn.Module):
+    def __init__(self, in_channels, out_channels, stride):
+        super().__init__()
+        self.layers = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(out_channels)
+        )
+        self.downsample = (
+            nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_channels)
+            ) if stride != 1 or in_channels != out_channels else nn.Identity()
+        )
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        residual = self.downsample(x)
+        x = self.layers(x)
+        return self.relu(x + residual)
+
+
+class AirNetExpert(nn.Module):
+    def __init__(self, in_shape: tuple, out_shape: tuple, prm: dict, device: torch.device) -> None:
+        super().__init__()
+        self.device = device
+        self.in_channels = in_shape[1]
+        self.image_size = in_shape[2]
+        self.num_classes = out_shape[0]
+        self.learning_rate = prm['lr']
+        self.momentum = prm['momentum']
+
+        channels = [64, 128, 256, 512]
+        init_block_channels = 64
+
+        self.features = self.build_features(init_block_channels, channels)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.classifier = nn.Linear(channels[-1], self.num_classes)
+
+    def build_features(self, init_block_channels, channels):
+        layers = [AirInitBlock(self.in_channels, init_block_channels)]
+        for i, out_channels in enumerate(channels):
+            layers.append(AirUnit(
+                in_channels=init_block_channels if i == 0 else channels[i - 1],
+                out_channels=out_channels,
+                stride=1 if i == 0 else 2))
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        return self.classifier(x)
+
+    def train_setup(self, prm):
+        self.to(self.device)
+        self.criteria = nn.CrossEntropyLoss().to(self.device)
+        self.optimizer = torch.optim.SGD(
+            self.parameters(),
+            lr=self.learning_rate,
+            momentum=self.momentum
+        )
+
+    def learn(self, train_data):
+        self.train()
+        for inputs, labels in train_data:
+            inputs, labels = inputs.to(self.device), labels.to(self.device)
+            self.optimizer.zero_grad()
+            outputs = self(inputs)
+            loss = self.criteria(outputs, labels)
+            loss.backward()
+            nn.utils.clip_grad_norm_(self.parameters(), 3)
+            self.optimizer.step()
+
+
+
+# ============================================================================
+# EXPERT 2: AlexNet
+# ============================================================================
+
+
+class AlexNetExpert(nn.Module):
+
+    def train_setup(self, prm):
+        self.to(self.device)
+        self.criteria = (nn.CrossEntropyLoss().to(self.device),)
+        self.optimizer = torch.optim.SGD(self.parameters(), lr=prm['lr'], momentum=prm['momentum'])
+
+    def learn(self, train_data):
+        for inputs, labels in train_data:
+            inputs, labels = inputs.to(self.device), labels.to(self.device)
+            self.optimizer.zero_grad()
+            outputs = self(inputs)
+            loss = self.criteria[0](outputs, labels)
+            loss.backward()
+            nn.utils.clip_grad_norm_(self.parameters(), 3)
+            self.optimizer.step()
+
+    def __init__(self, in_shape: tuple, out_shape: tuple, prm: dict, device: torch.device) -> None:
+        super().__init__()
+        self.device = device
+        self.features = nn.Sequential(
+            nn.Conv2d(in_shape[1], 64, kernel_size=11, stride=4, padding=2),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2),
+            nn.Conv2d(64, 192, kernel_size=5, padding=2),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2),
+            nn.Conv2d(192, 384, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(384, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2),
+        )
+        dropout: float = prm['dropout']
+        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
+        self.classifier = nn.Sequential(
+            nn.Dropout(p=dropout),
+            nn.Linear(256 * 6 * 6, 4096),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=dropout),
+            nn.Linear(4096, 4096),
+            nn.ReLU(inplace=True),
+            nn.Linear(4096, out_shape[0]),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.features(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.classifier(x)
+        return x
+
+
+
+# ============================================================================
+# EXPERT 3: DPN107
+# ============================================================================
+
+
+# Define DPNBlock with Group Convolutions
+class DPNBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride=1):
+        super(DPNBlock, self).__init__()
+        self.conv1 = nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=stride, padding=1, groups=4, bias=False
+        )
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(
+            out_channels, out_channels, kernel_size=3, padding=1, groups=4, bias=False
+        )
+        self.bn2 = nn.BatchNorm2d(out_channels)
+
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out += residual  # Residual connection
+        return self.relu(out)
+
+
+# Memory-Optimized DPN107
+class DPN107(nn.Module):
+    def __init__(self, in_channels, num_classes, num_blocks, growth_rate):
+        super(DPN107, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, growth_rate, kernel_size=3, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(growth_rate)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.blocks = nn.Sequential(
+            *[DPNBlock(growth_rate, growth_rate) for _ in range(num_blocks)]
+        )
+
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(growth_rate, num_classes)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.blocks(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+        return x
+
+
+class DPN107Expert(nn.Module):
+    def __init__(self, in_shape: tuple, out_shape: tuple, prm: dict, device: torch.device) -> None:
+        super(DPN107Expert, self).__init__()
+        self.device = device
+        model_class = DPN107
+        self.channel_number = in_shape[1]
+        self.image_size = in_shape[2]
+        self.class_number = out_shape[0]
+
+        self.model = model_class(self.channel_number, self.class_number, num_blocks=3, growth_rate=32)
+        self.learning_rate = prm['lr']
+        self.momentum = prm['momentum']
+
+    def forward(self, x):
+        return self.model(x)
+
+    def train_setup(self, prm):
+        self.to(self.device)
+        self.criteria = nn.CrossEntropyLoss().to(self.device)
+        self.optimizer = optim.SGD(self.parameters(), lr=self.learning_rate, momentum=self.momentum)
+
+    def learn(self, train_data):
+        self.train()
+        for inputs, labels in train_data:
+            inputs, labels = inputs.to(self.device), labels.to(self.device)
+            inputs = inputs.float()
+            self.optimizer.zero_grad()
+            outputs = self(inputs)
+            loss = self.criteria(outputs, labels)
+            loss.backward()
+            nn.utils.clip_grad_norm_(self.parameters(), 3)
+            self.optimizer.step()
+
+
+
+# ============================================================================
+# EXPERT 4: SqueezeNet-2
+# ============================================================================
+
+
+
+class Fire(nn.Module):
+    def __init__(self, inplanes: int, squeeze_planes: int, expand1x1_planes: int, expand3x3_planes: int) -> None:
+        super().__init__()
+        self.inplanes = inplanes
+        self.squeeze = nn.Conv2d(inplanes, squeeze_planes, kernel_size=1)
+        self.squeeze_activation = nn.ReLU(inplace=True)
+        self.expand1x1 = nn.Conv2d(squeeze_planes, expand1x1_planes, kernel_size=1)
+        self.expand1x1_activation = nn.ReLU(inplace=True)
+        self.expand3x3 = nn.Conv2d(squeeze_planes, expand3x3_planes, kernel_size=3, padding=1)
+        self.expand3x3_activation = nn.ReLU(inplace=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.squeeze_activation(self.squeeze(x))
+        return torch.cat(
+            [self.expand1x1_activation(self.expand1x1(x)), self.expand3x3_activation(self.expand3x3(x))], 1
+        )
+
+
+class SqueezeNet_2Expert(nn.Module):
+
+    def train_setup(self, prm):
+        self.to(self.device)
+        self.criteria = (nn.CrossEntropyLoss().to(self.device),)
+        self.optimizer = torch.optim.SGD(self.parameters(), lr=prm['lr'], momentum=prm['momentum'])
+
+    def learn(self, train_data):
+        for inputs, labels in train_data:
+            inputs, labels = inputs.to(self.device), labels.to(self.device)
+            self.optimizer.zero_grad()
+            outputs = self(inputs)
+            loss = self.criteria[0](outputs, labels)
+            loss.backward()
+            nn.utils.clip_grad_norm_(self.parameters(), 3)
+            self.optimizer.step()
+
+    def __init__(self, in_shape: tuple, out_shape: tuple, prm: dict, device: torch.device) -> None:
+        super().__init__()
+        self.device = device
+        version: str = "1_1"
+        num_classes: int = out_shape[0]
+        dropout: float = prm['dropout']
+        self.num_classes = num_classes
+        if version == "1_0":
+            self.features = nn.Sequential(
+                nn.Conv2d(in_shape[1], 96, kernel_size=7, stride=2),
+                nn.ReLU(inplace=True),
+                nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
+                Fire(96, 16, 64, 64),
+                Fire(128, 16, 64, 64),
+                Fire(128, 32, 128, 128),
+                nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
+                Fire(256, 32, 128, 128),
+                Fire(256, 48, 192, 192),
+                Fire(384, 48, 192, 192),
+                Fire(384, 64, 256, 256),
+                nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
+                Fire(512, 64, 256, 256),
+            )
+        elif version == "1_1":
+            self.features = nn.Sequential(
+                nn.Conv2d(in_shape[1], 64, kernel_size=3, stride=2),
+                nn.ReLU(inplace=True),
+                nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
+                Fire(64, 16, 64, 64),
+                Fire(128, 16, 64, 64),
+                nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
+                Fire(128, 32, 128, 128),
+                Fire(256, 32, 128, 128),
+                nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
+                Fire(256, 48, 192, 192),
+                Fire(384, 48, 192, 192),
+                Fire(384, 64, 256, 256),
+                Fire(512, 64, 256, 256),
+            )
+        else:
+            raise ValueError(f"Unsupported SqueezeNet version {version}: 1_0 or 1_1 expected")
+        final_conv = nn.Conv2d(512, self.num_classes, kernel_size=1)
+        self.classifier = nn.Sequential(
+            nn.Dropout(p=dropout), final_conv, nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d((1, 1))
+        )
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                if m is final_conv:
+                    init.normal_(m.weight, mean=0.0, std=0.01)
+                else:
+                    init.kaiming_uniform_(m.weight)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.features(x)
+        x = self.classifier(x)
+        return torch.flatten(x, 1)
+
+
+def _squeezenet(
+        version: str,
+        weights: Optional[WeightsEnum],
+        progress: bool,
+        **kwargs: Any,
+) -> SqueezeNet_2Expert:
+    if weights is not None:
+        _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"]))
+    model = SqueezeNet_2Expert(version, **kwargs)
+    if weights is not None:
+        model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True))
+
+    return model
+
+
+
+# ============================================================================
+# HETEROGENEOUS MOE GATE
+# ============================================================================
+class HeterogeneousGate(nn.Module):
+    """Lightweight CNN-based gating network that routes inputs to 4 experts."""
+    def __init__(self, input_channels, n_experts=4):
+        super().__init__()
+        self.n_experts = n_experts
+        self.gate_features = nn.Sequential(
+            nn.Conv2d(input_channels, 64, 3, stride=2, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, 128, 3, stride=2, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+            nn.AdaptiveAvgPool2d((1, 1)),
+        )
+        self.gate = nn.Sequential(
+            nn.Linear(128, 64),
+            nn.ReLU(),
+            nn.Dropout(0.1),
+            nn.Linear(64, n_experts),
+        )
+        self.temperature = nn.Parameter(torch.ones(1) * 2.0)
+
+    def forward(self, x):
+        features = self.gate_features(x).flatten(1)
+        logits = self.gate(features) / torch.clamp(self.temperature, 0.5, 5.0)
+        if self.training:
+            logits = logits + torch.randn_like(logits) * 0.1
+        return F.softmax(logits, dim=-1), logits
+
+
+# ============================================================================
+# HETEROGENEOUS MOE NET — AirNet + AlexNet + DPN107 + SqueezeNet_2
+# ============================================================================
+class Net(nn.Module):
+    def __init__(self, in_shape: tuple, out_shape: tuple, prm: dict, device: torch.device) -> None:
+        super().__init__()
+        self.device = device
+        n_experts = 4
+
+        # Defaults for expert-specific hyperparameters
+        pass  # no extra defaults needed
+        self.experts = nn.ModuleList([
+            AirNetExpert(in_shape, out_shape, prm, device),
+            AlexNetExpert(in_shape, out_shape, prm, device),
+            DPN107Expert(in_shape, out_shape, prm, device),
+            SqueezeNet_2Expert(in_shape, out_shape, prm, device),
+        ])
+        self.gate = HeterogeneousGate(in_shape[1], n_experts=n_experts)
+
+        self.load_balance_weight = 0.01
+        self.label_smoothing = 0.1
+        self.mixup_alpha = 0.2
+        self.to(device)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate_weights, _ = self.gate(x)
+        expert_outputs = torch.stack([e(x) for e in self.experts], dim=2)
+        return torch.sum(expert_outputs * gate_weights.unsqueeze(1), dim=2)
+
+    def _load_balance_loss(self):
+        if not hasattr(self, '_last_gw'):
+            return torch.tensor(0.0, device=self.device)
+        gw = self._last_gw
+        usage = gw.sum(dim=0)
+        target = gw.sum() / len(self.experts)
+        return F.mse_loss(usage, target.expand_as(usage))
+
+    def _mixup_data(self, x, y):
+        lam = np.random.beta(self.mixup_alpha, self.mixup_alpha) if self.mixup_alpha > 0 else 1.0
+        idx = torch.randperm(x.size(0), device=self.device)
+        return lam * x + (1 - lam) * x[idx], y, y[idx], lam
+
+    def train_setup(self, prm):
+        self.to(self.device)
+        self.criteria = nn.CrossEntropyLoss(label_smoothing=self.label_smoothing)
+        expert_params, gate_params = [], list(self.gate.parameters())
+        for e in self.experts:
+            expert_params.extend(list(e.parameters()))
+        self.optimizer = torch.optim.AdamW([
+            {'params': expert_params, 'lr': prm.get('lr', 0.001), 'weight_decay': 5e-4},
+            {'params': gate_params, 'lr': prm.get('lr', 0.001) * 2, 'weight_decay': 2.5e-4},
+        ])
+        self.warmup_scheduler = torch.optim.lr_scheduler.LinearLR(
+            self.optimizer, start_factor=0.1, total_iters=5
+        )
+        self.main_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            self.optimizer, T_max=50, eta_min=1e-6
+        )
+        self._current_epoch = 0
+
+    def learn(self, train_data):
+        self.train()
+        for inputs, labels in train_data:
+            inputs = inputs.to(self.device, dtype=torch.float32)
+            labels = labels.to(self.device, dtype=torch.long)
+            if self.mixup_alpha > 0:
+                inputs, labels_a, labels_b, lam = self._mixup_data(inputs, labels)
+                self.optimizer.zero_grad()
+                outputs = self(inputs)
+                loss = lam * self.criteria(outputs, labels_a) + (1 - lam) * self.criteria(outputs, labels_b)
+            else:
+                self.optimizer.zero_grad()
+                outputs = self(inputs)
+                loss = self.criteria(outputs, labels)
+            self._last_gw = self.gate(inputs)[0].detach()
+            loss = loss + self.load_balance_weight * self._load_balance_loss()
+            loss.backward()
+            nn.utils.clip_grad_norm_(self.parameters(), 1.0)
+            self.optimizer.step()
+        self._current_epoch += 1
+        if self._current_epoch <= 5:
+            self.warmup_scheduler.step()
+        else:
+            self.main_scheduler.step()
+

test/nn/MoE4Own-batch-00057-20260408-085045-459eacce1d6ebe52380edb28feb88826.py

@@ -0,0 +1,532 @@
+# Auto-generated 4-Expert Heterogeneous MoE: AirNet + BagNet + DPN68 + DarkNet
+# Four-expert mixture with learned gating, mixup, and LR scheduling
+
+import numpy as np
+import torch.nn.functional as F
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+
+def supported_hyperparameters():
+    return {'dropout', 'lr', 'momentum'}
+
+
+
+# ============================================================================
+# EXPERT 1: AirNet
+# ============================================================================
+
+
+class AirInitBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.layers = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class AirUnit(nn.Module):
+    def __init__(self, in_channels, out_channels, stride):
+        super().__init__()
+        self.layers = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(out_channels)
+        )
+        self.downsample = (
+            nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_channels)
+            ) if stride != 1 or in_channels != out_channels else nn.Identity()
+        )
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        residual = self.downsample(x)
+        x = self.layers(x)
+        return self.relu(x + residual)
+
+
+class AirNetExpert(nn.Module):
+    def __init__(self, in_shape: tuple, out_shape: tuple, prm: dict, device: torch.device) -> None:
+        super().__init__()
+        self.device = device
+        self.in_channels = in_shape[1]
+        self.image_size = in_shape[2]
+        self.num_classes = out_shape[0]
+        self.learning_rate = prm['lr']
+        self.momentum = prm['momentum']
+
+        channels = [64, 128, 256, 512]
+        init_block_channels = 64
+
+        self.features = self.build_features(init_block_channels, channels)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.classifier = nn.Linear(channels[-1], self.num_classes)
+
+    def build_features(self, init_block_channels, channels):
+        layers = [AirInitBlock(self.in_channels, init_block_channels)]
+        for i, out_channels in enumerate(channels):
+            layers.append(AirUnit(
+                in_channels=init_block_channels if i == 0 else channels[i - 1],
+                out_channels=out_channels,
+                stride=1 if i == 0 else 2))
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        return self.classifier(x)
+
+    def train_setup(self, prm):
+        self.to(self.device)
+        self.criteria = nn.CrossEntropyLoss().to(self.device)
+        self.optimizer = torch.optim.SGD(
+            self.parameters(),
+            lr=self.learning_rate,
+            momentum=self.momentum
+        )
+
+    def learn(self, train_data):
+        self.train()
+        for inputs, labels in train_data:
+            inputs, labels = inputs.to(self.device), labels.to(self.device)
+            self.optimizer.zero_grad()
+            outputs = self(inputs)
+            loss = self.criteria(outputs, labels)
+            loss.backward()
+            nn.utils.clip_grad_norm_(self.parameters(), 3)
+            self.optimizer.step()
+
+
+
+# ============================================================================
+# EXPERT 2: BagNet
+# ============================================================================
+
+
+class BagNetBottleneck(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, bottleneck_factor=4):
+        super().__init__()
+        mid_channels = out_channels // bottleneck_factor
+
+        self.conv1 = self.conv1x1_block(in_channels, mid_channels)
+        self.conv2 = self.conv_block(mid_channels, mid_channels, kernel_size, stride)
+        self.conv3 = self.conv1x1_block(mid_channels, out_channels, activation=False)
+
+    def conv1x1_block(self, in_channels, out_channels, activation=True):
+        layers = [nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)]
+        if activation:
+            layers.append(nn.ReLU(inplace=True))
+        return nn.Sequential(*layers)
+
+    def conv_block(self, in_channels, out_channels, kernel_size, stride):
+        padding = (kernel_size - 1) // 2
+        return nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        return x
+
+
+class BagNetUnit(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride):
+        super().__init__()
+        self.resize_identity = (in_channels != out_channels) or (stride != 1)
+        self.body = BagNetBottleneck(in_channels, out_channels, kernel_size, stride)
+
+        if self.resize_identity:
+            self.identity_conv = self.conv1x1_block(in_channels, out_channels, activation=False)
+        self.activ = nn.ReLU(inplace=True)
+
+    def conv1x1_block(self, in_channels, out_channels, activation=True):
+        layers = [nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)]
+        if activation:
+            layers.append(nn.ReLU(inplace=True))
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        identity = x
+        if self.resize_identity:
+            identity = self.identity_conv(x)
+
+        x = self.body(x)
+
+        if x.size(2) != identity.size(2) or x.size(3) != identity.size(3):
+            identity = nn.functional.interpolate(identity, size=(x.size(2), x.size(3)), mode='bilinear', align_corners=False)
+
+        return self.activ(x + identity)
+
+
+class BagNetExpert(nn.Module):
+    def __init__(self, in_shape: tuple, out_shape: tuple, prm: dict, device: torch.device) -> None:
+        super().__init__()
+        self.device = device
+        channel_number = in_shape[1]
+        image_size = in_shape[2]
+        class_number = out_shape[0]
+        learning_rate = prm['lr']
+        momentum = prm['momentum']
+        dropout = prm['dropout']
+
+        self.channels = [[64, 64, 64], [128, 128, 128], [256, 256, 256], [512, 512, 512]]
+        self.in_size = image_size
+        self.num_classes = class_number
+
+        self.features = nn.Sequential(
+            nn.Conv2d(channel_number, 64, kernel_size=7, stride=2, padding=3, bias=False),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
+        )
+
+        in_channels = 64
+        for i, stage_channels in enumerate(self.channels):
+            stage = nn.Sequential()
+            for j, out_channels in enumerate(stage_channels):
+                stride = 2 if (j == 0 and i > 0) else 1
+                stage.add_module(f"unit{j + 1}", BagNetUnit(in_channels, out_channels, kernel_size=3, stride=stride))
+                in_channels = out_channels
+            self.features.add_module(f"stage{i + 1}", stage)
+
+        self.features.add_module("final_pool", nn.AdaptiveAvgPool2d(1))
+        self.output = nn.Linear(in_channels, self.num_classes)
+
+        self.learning_rate = learning_rate
+        self.momentum = momentum
+        self.dropout = dropout
+
+    def forward(self, x):
+        x = self.features(x)
+        x = torch.flatten(x, 1)
+        return self.output(x)
+
+    def train_setup(self, prm):
+        self.to(self.device)
+        self.criteria = nn.CrossEntropyLoss().to(self.device)
+        self.optimizer = torch.optim.SGD(self.parameters(), lr=prm['lr'], momentum=prm['momentum'],)
+
+        if self.dropout > 0:
+            self.dropout_layer = nn.Dropout(self.dropout)
+
+    def learn(self, train_data):
+        self.train()
+        for inputs, labels in train_data:
+            inputs, labels = inputs.to(self.device), labels.to(self.device)
+            self.optimizer.zero_grad()
+            outputs = self(inputs)
+            loss = self.criteria(outputs, labels)
+            loss.backward()
+            nn.utils.clip_grad_norm_(self.parameters(), 3)
+            self.optimizer.step()
+
+
+
+# ============================================================================
+# EXPERT 3: DPN68
+# ============================================================================
+
+
+class DPNBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(DPNBlock, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = out + residual
+        return self.relu(out)
+
+
+class DPN68(nn.Module):
+    def __init__(self, in_channels, num_classes, num_blocks, growth_rate):
+        super(DPN68, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, growth_rate, kernel_size=3, stride=1, padding=1)
+        self.bn1 = nn.BatchNorm2d(growth_rate)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.blocks = nn.Sequential(
+            *[DPNBlock(growth_rate, growth_rate) for _ in range(num_blocks)]
+        )
+
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(growth_rate, num_classes)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.blocks(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+        return x
+
+
+class DPN68Expert(nn.Module):
+    def __init__(self, in_shape: tuple, out_shape: tuple, prm: dict, device: torch.device) -> None:
+        super(DPN68Expert, self).__init__()
+        self.device = device
+        model_class = DPN68
+        self.channel_number = in_shape[1]
+        self.image_size = in_shape[2]
+        self.class_number = out_shape[0]
+        self.model = model_class(self.channel_number, self.class_number, num_blocks=3, growth_rate=32)
+
+    def forward(self, x):
+        return self.model(x)
+
+    def train_setup(self, prm):
+        self.to(self.device)
+        self.criteria = nn.CrossEntropyLoss().to(self.device)
+        self.optimizer = optim.SGD(self.parameters(), lr=prm['lr'], momentum=prm['momentum'])
+
+    def learn(self, train_data):
+        self.train()
+        for inputs, labels in train_data:
+            inputs, labels = inputs.to(self.device), labels.to(self.device)
+            self.optimizer.zero_grad()
+            outputs = self(inputs)
+            loss = self.criteria(outputs, labels)
+            loss.backward()
+            nn.utils.clip_grad_norm_(self.parameters(), 3)
+            self.optimizer.step()
+
+
+
+# ============================================================================
+# EXPERT 4: DarkNet
+# ============================================================================
+
+
+class DarkNetUnit(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, pointwise: bool, alpha: float):
+        super(DarkNetUnit, self).__init__()
+        self.activation = nn.LeakyReLU(negative_slope=alpha, inplace=True)
+        if pointwise:
+            self.conv = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, bias=False),
+                nn.BatchNorm2d(out_channels),
+                self.activation
+            )
+        else:
+            self.conv = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False),
+                nn.BatchNorm2d(out_channels),
+                self.activation
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)
+
+
+class DarkNetExpert(nn.Module):
+    def __init__(self, in_shape: tuple, out_shape: tuple, prm: dict, device: torch.device) -> None:
+        super(DarkNetExpert, self).__init__()
+        self.device = device
+        channels: list = None
+        odd_pointwise: bool = True
+        alpha: float = 0.1
+        in_channels = in_shape[1]
+        image_size = in_shape[2]
+        num_classes = out_shape[0]
+
+        if channels is None:
+            channels = [[64, 64, 64], [128, 128, 128], [256, 256, 256], [512, 512, 512]]
+
+        self.features = nn.Sequential()
+        for i, channels_per_stage in enumerate(channels):
+            stage = nn.Sequential()
+            for j, out_channels in enumerate(channels_per_stage):
+                pointwise = (len(channels_per_stage) > 1) and not (((j + 1) % 2 == 1) ^ odd_pointwise)
+                stage.add_module(f"unit{j + 1}", DarkNetUnit(in_channels, out_channels, pointwise, alpha))
+                in_channels = out_channels
+            if i != len(channels) - 1:
+                stage.add_module(f"pool{i + 1}", nn.MaxPool2d(kernel_size=2, stride=2))
+            self.features.add_module(f"stage{i + 1}", stage)
+
+        final_feature_map_size = image_size // (2 ** (len(channels) - 1))
+
+        self.output = nn.Sequential(
+            nn.Conv2d(in_channels=in_channels, out_channels=num_classes, kernel_size=1),
+            nn.LeakyReLU(negative_slope=alpha, inplace=True),
+            nn.AdaptiveAvgPool2d(output_size=(1, 1))
+        )
+
+        self._initialize_weights()
+
+    def _initialize_weights(self):
+        for module in self.modules():
+            if isinstance(module, nn.Conv2d):
+                nn.init.kaiming_uniform_(module.weight, mode='fan_in', nonlinearity='leaky_relu')
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.features(x)
+        x = self.output(x)
+        x = x.view(x.size(0), -1)
+        return x
+
+    def train_setup(self, prm: dict):
+        self.to(self.device)
+        learning_rate = float(prm.get("lr", 0.01))
+        momentum = float(prm.get("momentum", 0.9))
+        self.criteria = nn.CrossEntropyLoss()
+        self.optimizer = optim.SGD(self.parameters(), lr=learning_rate, momentum=momentum)
+        self.to(self.device)
+
+    def learn(self, train_data: torch.utils.data.DataLoader):
+        self.train()
+        for inputs, targets in train_data:
+            inputs, targets = inputs.to(next(self.parameters()).device), targets.to(next(self.parameters()).device)
+            self.optimizer.zero_grad()
+            outputs = self(inputs)
+            loss = self.criteria(outputs, targets)
+            loss.backward()
+            self.optimizer.step()
+
+
+
+# ============================================================================
+# HETEROGENEOUS MOE GATE
+# ============================================================================
+class HeterogeneousGate(nn.Module):
+    """Lightweight CNN-based gating network that routes inputs to 4 experts."""
+    def __init__(self, input_channels, n_experts=4):
+        super().__init__()
+        self.n_experts = n_experts
+        self.gate_features = nn.Sequential(
+            nn.Conv2d(input_channels, 64, 3, stride=2, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, 128, 3, stride=2, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+            nn.AdaptiveAvgPool2d((1, 1)),
+        )
+        self.gate = nn.Sequential(
+            nn.Linear(128, 64),
+            nn.ReLU(),
+            nn.Dropout(0.1),
+            nn.Linear(64, n_experts),
+        )
+        self.temperature = nn.Parameter(torch.ones(1) * 2.0)
+
+    def forward(self, x):
+        features = self.gate_features(x).flatten(1)
+        logits = self.gate(features) / torch.clamp(self.temperature, 0.5, 5.0)
+        if self.training:
+            logits = logits + torch.randn_like(logits) * 0.1
+        return F.softmax(logits, dim=-1), logits
+
+
+# ============================================================================
+# HETEROGENEOUS MOE NET — AirNet + BagNet + DPN68 + DarkNet
+# ============================================================================
+class Net(nn.Module):
+    def __init__(self, in_shape: tuple, out_shape: tuple, prm: dict, device: torch.device) -> None:
+        super().__init__()
+        self.device = device
+        n_experts = 4
+
+        # Defaults for expert-specific hyperparameters
+        pass  # no extra defaults needed
+        self.experts = nn.ModuleList([
+            AirNetExpert(in_shape, out_shape, prm, device),
+            BagNetExpert(in_shape, out_shape, prm, device),
+            DPN68Expert(in_shape, out_shape, prm, device),
+            DarkNetExpert(in_shape, out_shape, prm, device),
+        ])
+        self.gate = HeterogeneousGate(in_shape[1], n_experts=n_experts)
+
+        self.load_balance_weight = 0.01
+        self.label_smoothing = 0.1
+        self.mixup_alpha = 0.2
+        self.to(device)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate_weights, _ = self.gate(x)
+        expert_outputs = torch.stack([e(x) for e in self.experts], dim=2)
+        return torch.sum(expert_outputs * gate_weights.unsqueeze(1), dim=2)
+
+    def _load_balance_loss(self):
+        if not hasattr(self, '_last_gw'):
+            return torch.tensor(0.0, device=self.device)
+        gw = self._last_gw
+        usage = gw.sum(dim=0)
+        target = gw.sum() / len(self.experts)
+        return F.mse_loss(usage, target.expand_as(usage))
+
+    def _mixup_data(self, x, y):
+        lam = np.random.beta(self.mixup_alpha, self.mixup_alpha) if self.mixup_alpha > 0 else 1.0
+        idx = torch.randperm(x.size(0), device=self.device)
+        return lam * x + (1 - lam) * x[idx], y, y[idx], lam
+
+    def train_setup(self, prm):
+        self.to(self.device)
+        self.criteria = nn.CrossEntropyLoss(label_smoothing=self.label_smoothing)
+        expert_params, gate_params = [], list(self.gate.parameters())
+        for e in self.experts:
+            expert_params.extend(list(e.parameters()))
+        self.optimizer = torch.optim.AdamW([
+            {'params': expert_params, 'lr': prm.get('lr', 0.001), 'weight_decay': 5e-4},
+            {'params': gate_params, 'lr': prm.get('lr', 0.001) * 2, 'weight_decay': 2.5e-4},
+        ])
+        self.warmup_scheduler = torch.optim.lr_scheduler.LinearLR(
+            self.optimizer, start_factor=0.1, total_iters=5
+        )
+        self.main_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            self.optimizer, T_max=50, eta_min=1e-6
+        )
+        self._current_epoch = 0
+
+    def learn(self, train_data):
+        self.train()
+        for inputs, labels in train_data:
+            inputs = inputs.to(self.device, dtype=torch.float32)
+            labels = labels.to(self.device, dtype=torch.long)
+            if self.mixup_alpha > 0:
+                inputs, labels_a, labels_b, lam = self._mixup_data(inputs, labels)
+                self.optimizer.zero_grad()
+                outputs = self(inputs)
+                loss = lam * self.criteria(outputs, labels_a) + (1 - lam) * self.criteria(outputs, labels_b)
+            else:
+                self.optimizer.zero_grad()
+                outputs = self(inputs)
+                loss = self.criteria(outputs, labels)
+            self._last_gw = self.gate(inputs)[0].detach()
+            loss = loss + self.load_balance_weight * self._load_balance_loss()
+            loss.backward()
+            nn.utils.clip_grad_norm_(self.parameters(), 1.0)
+            self.optimizer.step()
+        self._current_epoch += 1
+        if self._current_epoch <= 5:
+            self.warmup_scheduler.step()
+        else:
+            self.main_scheduler.step()
+

test/nn/alt-nn1-f49a24e1253f3039135ce0fdebed8697.py

@@ -0,0 +1,89 @@
+import torch
+import torch.nn as nn
+from torch.nn import GELU
+
+def supported_hyperparameters():
+    return {'lr', 'momentum', 'dropout'}
+
+
+class Net(nn.Module):
+    def train_setup(self, prm):
+        self.to(self.device)
+        self.criteria = (nn.CrossEntropyLoss().to(self.device),)
+        self.optimizer = torch.optim.SGD(
+            self.parameters(),
+            lr=prm['lr'],
+            momentum=prm['momentum']
+        )
+
+    def learn(self, train_data):
+        self.train()
+        for inputs, labels in train_data:
+            inputs, labels = inputs.to(self.device), labels.to(self.device)
+            self.optimizer.zero_grad()
+            outputs = self(inputs)
+            loss = self.criteria[0](outputs, labels)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(self.parameters(), 3)
+            self.optimizer.step()
+
+    def __init__(self, in_shape: tuple, out_shape: tuple, prm: dict, device: torch.device) -> None:
+        super().__init__()
+        self.device = device
+        layers = []
+        in_channels = in_shape[1]
+
+        layers += [
+            nn.Conv2d(in_channels, 96, kernel_size=7,
+                      stride=3, padding=2),
+            GELU(),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+        ]
+        in_channels = 96
+
+        layers += [
+            nn.Conv2d(in_channels, 192, kernel_size=3, padding=2),
+            GELU(),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+        ]
+        in_channels = 192
+
+        layers += [
+            nn.Conv2d(in_channels, 384, kernel_size=3, padding=1),
+            GELU(),
+        ]
+        in_channels = 384
+
+        layers += [
+            nn.Conv2d(in_channels, 256, kernel_size=3, padding=1),
+            GELU(),
+        ]
+        in_channels = 256
+
+        layers += [
+            nn.Conv2d(in_channels, 256, kernel_size=3, padding=1),
+            GELU(),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+        ]
+        in_channels = 256
+        self.features = nn.Sequential(*layers)
+        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
+
+        dropout_p = prm['dropout']
+        classifier_input_features = in_channels * 6 * 6
+        self.classifier = nn.Sequential(
+            nn.Dropout(p=dropout_p),
+            nn.Linear(classifier_input_features, 4096),
+            GELU(),
+            nn.Dropout(p=dropout_p),
+            nn.Linear(4096, 2048),
+            GELU(),
+            nn.Linear(2048, out_shape[0]),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.features(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.classifier(x)
+        return x

test/nn/alt-nn2-54f7346afec7cd2c53b2f5654e8c7da9.py

@@ -0,0 +1,92 @@
+import torch
+import torch.nn as nn
+
+def supported_hyperparameters():
+    return {'lr', 'momentum', 'dropout'}
+
+class Net(nn.Module):
+
+    def train_setup(self, prm):
+        self.to(self.device)
+        self.criteria = (nn.CrossEntropyLoss().to(self.device),)
+        self.optimizer = torch.optim.SGD(
+            self.parameters(),
+            lr=prm['lr'],
+            momentum=prm['momentum']
+        )
+
+    def learn(self, train_data):
+        self.train()
+        for inputs, labels in train_data:
+            inputs, labels = inputs.to(self.device), labels.to(self.device)
+            self.optimizer.zero_grad()
+            outputs = self(inputs)
+            loss = self.criteria[0](outputs, labels)
+            loss.backward()
+            nn.utils.clip_grad_norm_(self.parameters(), 3)
+            self.optimizer.step()
+
+    def __init__(self, in_shape: tuple, out_shape: tuple, prm: dict, device: torch.device) -> None:
+        super().__init__()
+        self.device = device
+        layers = []
+        in_channels = in_shape[1]
+
+        layers += [
+            nn.Conv2d(in_channels, 64, kernel_size=7, stride=3, padding=2),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        ]
+        in_channels = 64
+
+        layers += [
+            nn.Conv2d(in_channels, 192, kernel_size=3, padding=2),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        ]
+        in_channels = 192
+
+        layers += [
+            nn.Conv2d(in_channels, 384, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+        ]
+        in_channels = 384
+
+        layers += [
+            nn.Conv2d(in_channels, 384, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+        ]
+        in_channels = 384
+
+        layers += [
+            nn.Conv2d(in_channels, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        ]
+        in_channels = 256
+
+        self.features = nn.Sequential(*layers)
+        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
+
+        dropout_p = prm['dropout']
+        classifier_input_features = in_channels * 6 * 6
+        self.classifier = nn.Sequential(
+            nn.Dropout(p=dropout_p),
+            nn.Linear(classifier_input_features, 3072),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=dropout_p),
+            nn.Linear(3072, 3072),
+            nn.ReLU(inplace=True),
+            nn.Linear(3072, out_shape[0])
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.features(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.classifier(x)
+        return x
+
+                                         
+                                                                                                                                                                                                                                                                            
+addon_accuracy: 0.7753

test/nn/ga-524.py

@@ -0,0 +1,92 @@
+import torch
+import torch.nn as nn
+
+
+def supported_hyperparameters():
+    return {'lr', 'momentum', 'dropout'}
+
+
+class Net(nn.Module):
+    def train_setup(self, prm):
+        self.to(self.device)
+        self.criteria = (nn.CrossEntropyLoss().to(self.device),)
+        self.optimizer = torch.optim.SGD(
+            self.parameters(),
+            lr=prm['lr'],
+            momentum=prm['momentum']
+        )
+
+    def learn(self, train_data):
+        self.train()
+        for inputs, labels in train_data:
+            inputs, labels = inputs.to(self.device), labels.to(self.device)
+            self.optimizer.zero_grad()
+            outputs = self(inputs)
+            loss = self.criteria[0](outputs, labels)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(self.parameters(), 3)
+            self.optimizer.step()
+
+    def __init__(self, in_shape: tuple, out_shape: tuple, prm: dict, device: torch.device) -> None:
+        super().__init__()
+        self.device = device
+        layers = []
+        in_channels = in_shape[1]
+
+        layers += [
+            nn.Conv2d(in_channels, 64, kernel_size=7,
+                      stride=3, padding=2),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+        ]
+        in_channels = 64
+
+        layers += [
+            nn.Conv2d(in_channels, 192, kernel_size=3, padding=2),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+        ]
+        in_channels = 192
+
+        layers += [
+            nn.Conv2d(in_channels, 440, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+        ]
+        in_channels = 440
+
+        layers += [
+            nn.Conv2d(in_channels, 384, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+        ]
+        in_channels = 384
+
+        layers += [
+            nn.Conv2d(in_channels, 192, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+        ]
+        in_channels = 192
+        self.features = nn.Sequential(*layers)
+        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
+
+        dropout_p = prm['dropout']
+        classifier_input_features = in_channels * 6 * 6
+        self.classifier = nn.Sequential(
+            nn.Dropout(p=dropout_p),
+            nn.Linear(classifier_input_features, 2048),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=dropout_p),
+            nn.Linear(2048, 4096),
+            nn.ReLU(inplace=True),
+            nn.Linear(4096, out_shape[0]),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.features(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.classifier(x)
+        return x
+
+# --- Chromosome used to generate this model ---
+# Chromosome: {'conv1_filters': 64, 'conv1_kernel': 7, 'conv1_stride': 3, 'conv2_filters': 192, 'conv2_kernel': 3, 'conv3_filters': 440, 'conv4_filters': 384, 'conv5_filters': 192, 'fc1_neurons': 2048, 'fc2_neurons': 4096, 'lr': 0.01, 'momentum': 0.95, 'dropout': 0.4}

test/nn/rl-bb-init-b4de916c853aa8113c079065856dc78a.py

@@ -0,0 +1,210 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import gc
+import torchvision
+from torch.nn import MaxPool2d
+from torch.amp import autocast, GradScaler
+
+# -------------------------------------------------
+# Wrapper & Helpers
+# -------------------------------------------------
+class TorchVision(nn.Module):
+    def __init__(self, model: str, weights: str = "DEFAULT", unwrap: bool = True, truncate: int = 1, in_channels: int = 3):
+        super().__init__()
+        self.adapter = nn.Conv2d(in_channels, 3, kernel_size=1) if in_channels != 3 else nn.Identity()
+        if hasattr(torchvision.models, "get_model"):
+            self.m = torchvision.models.get_model(model, weights=weights)
+        else:
+            self.m = torchvision.models.__dict__[model](pretrained=bool(weights))
+        if unwrap:
+            layers = []
+            for name, module in self.m.named_children():
+                if "aux" in name.lower(): continue
+                layers.append(module)
+            self.m = nn.Sequential(*(layers[:-truncate] if truncate else layers))
+        else:
+            self.m.head = nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.dim() == 2:
+            x = x.unsqueeze(-1).unsqueeze(-1)
+        return self.m(self.adapter(x))
+
+def adaptive_pool_flatten(x):
+    if x.ndim == 4: return torch.nn.functional.adaptive_avg_pool2d(x, (1, 1)).flatten(1)
+    if x.ndim == 3: return x.mean(dim=1)
+    return x.flatten(1) if x.ndim > 2 else x
+
+# -------------------------------------------------
+# AMP / Hyperparams
+# -------------------------------------------------
+from torch.amp import autocast, GradScaler
+def autocast_ctx(enabled=True):
+    return autocast("cuda", enabled=enabled)
+def make_scaler(enabled=True):
+    return GradScaler("cuda", enabled=enabled)
+
+def supported_hyperparameters():
+    return { 'lr', 'dropout', 'momentum' }
+
+# -------------------------------------------------
+# Generated Conv Block (Crucial Fix)
+# -------------------------------------------------
+def drop_conv3x3_block(in_channels, out_channels, stride=1, padding=1, bias=False, dropout_prob=0.0):
+    return nn.Sequential(
+    nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=padding, bias=bias),
+        nn.GELU(),
+        nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=bias),
+        nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=bias),
+        nn.Dropout2d(p=dropout_prob) if dropout_prob > 0 else nn.Identity()
+    )
+
+# -------------------------------------------------
+# Fractal Components (Using drop_conv3x3_block)
+# -------------------------------------------------
+class FractalBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, num_columns, loc_drop_prob, dropout_prob):
+        super().__init__()
+        self.num_columns = int(num_columns)
+        depth = 2 ** max(self.num_columns - 1, 0)
+        blocks = []
+        for i in range(depth):
+            level = nn.ModuleList()
+            for j in range(self.num_columns):
+                if (i + 1) % (2 ** j) == 0:
+                    in_ch_ij = in_channels if (i + 1 == 2 ** j) else out_channels
+                    level.append(drop_conv3x3_block(in_ch_ij, out_channels, dropout_prob=dropout_prob))
+            blocks.append(level)
+        self.blocks = nn.ModuleList(blocks)
+        self.use_checkpoint_per_subblock = False
+
+    def forward(self, x):
+        outs = [x] * self.num_columns
+        for level_block in self.blocks:
+            outs_i = [blk(inp) for blk, inp in zip(level_block, outs)]
+            joined = torch.stack(outs_i, dim=0).mean(dim=0)
+            outs[:len(level_block)] = [joined] * len(level_block)
+        return outs[0]
+
+class FractalUnit(nn.Module):
+    def __init__(self, in_channels, out_channels, num_columns, loc_drop_prob, dropout_prob):
+        super().__init__()
+        self.block = FractalBlock(in_channels, out_channels, num_columns, loc_drop_prob, dropout_prob)
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+
+    def forward(self, x):
+        return self.pool(self.block(x))
+
+# -------------------------------------------------
+# Net (Dynamic Fusion Architecture)
+# -------------------------------------------------
+class Net(nn.Module):
+    def __init__(self, in_shape: tuple, out_shape: tuple, prm: dict, device: torch.device) -> None:
+        super().__init__()
+        self.device = device
+        self.pattern = "Split_A_Parallel_BF"
+        self.use_amp = prm.get("use_amp", False)
+        dropout_prob = float(prm.get('dropout', 0.1))
+
+        # --- Filled by Generator ---
+        self.backbone_a = TorchVision("mnasnet1_0", in_channels=3).to(device)
+
+        self.features = nn.Sequential()
+        curr_ch = 912
+        channels = [64 * (2 ** i) for i in range(1)]
+        for i, out_ch in enumerate(channels):
+            self.features.add_module(f"unit{i+1}", FractalUnit(curr_ch, out_ch, 3, 0.15, dropout_prob))
+            curr_ch = out_ch
+
+        self.backbone_b = TorchVision("regnet_x_1_6gf", in_channels=3).to(device)
+
+        self.infer_dimensions_dynamically(in_shape, out_shape[0])
+        self._scaler = GradScaler("cuda", enabled=self.use_amp)
+
+    def infer_dimensions_dynamically(self, in_shape, num_classes):
+        self.to(self.device)
+        self.eval()
+        with torch.no_grad():
+            C = in_shape[1] if len(in_shape)==4 else in_shape[0]
+            dummy = torch.zeros(1, C, 224, 224).to(self.device)
+            output_feat = self.forward(dummy, is_probing=True)
+            dim_fused = output_feat.shape[1]
+
+        self.classifier = nn.Linear(dim_fused, num_classes)
+        self.train()
+
+    def _init_params(self):
+        for m in self.features.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_uniform_(m.weight, a=np.sqrt(5))
+                if m.bias is not None: nn.init.zeros_(m.bias)
+
+    @staticmethod
+    def _norm4d(x: torch.Tensor) -> torch.Tensor:
+        if x.dim() == 4: return x
+        if x.dim() == 5:
+            B, T, C, H, W = x.shape
+            return x.reshape(B * T, C, H, W)
+        raise ValueError(f"Expected 4D/5D input, got {tuple(x.shape)}")
+
+    
+    def forward(self, x: torch.Tensor, is_probing: bool = False) -> torch.Tensor:
+        x = self._norm4d(x).to(self.device)
+        
+        f_a = adaptive_pool_flatten(self.backbone_a(x))
+        
+        x_bf = self.backbone_b(x)
+        if x_bf.dim() == 2:
+            x_bf = x_bf.unsqueeze(-1).unsqueeze(-1)
+        if x_bf.shape[-1] < 14:
+            x_bf = torch.nn.functional.interpolate(x_bf, size=(14,14), mode='nearest')
+            
+        f_bf = adaptive_pool_flatten(self.features(x_bf))
+        fused = torch.cat([f_a, f_bf], dim=1)
+        if is_probing: return fused
+        return self.classifier(fused)
+    
+
+    def train_setup(self, prm):
+        self.to(self.device)
+        self.criterion = nn.CrossEntropyLoss().to(self.device)
+        self.optimizer = torch.optim.SGD(
+            self.parameters(),
+            lr=prm['lr'],
+            momentum=prm['momentum']
+        )
+        self._scaler = make_scaler(enabled=self.use_amp)
+
+    def learn(self, train_data):
+        self.train()
+        scaler = self._scaler
+        train_iter = iter(train_data)
+        try:
+            for batch_idx, (inputs, labels) in enumerate(train_iter):
+                inputs = inputs.to(self.device).float()
+                labels = labels.to(self.device)
+                self.optimizer.zero_grad(set_to_none=True)
+
+                with autocast_ctx(enabled=self.use_amp):
+                    outputs = self(inputs)
+                    loss = self.criterion(outputs, labels)
+
+                if not torch.isfinite(loss):
+                    continue
+
+                if self.use_amp:
+                    scaler.scale(loss).backward()
+                    scaler.unscale_(self.optimizer)
+                    nn.utils.clip_grad_norm_(self.parameters(), 3.0)
+                    scaler.step(self.optimizer)
+                    scaler.update()
+                else:
+                    loss.backward()
+                    nn.utils.clip_grad_norm_(self.parameters(), 3.0)
+                    self.optimizer.step()
+        finally:
+            if hasattr(train_iter, 'shutdown'):
+                train_iter.shutdown()
+            del train_iter
+            gc.collect()

test/nn/unq-03221882bffefdd92eb8ae1d1fb46b22.py

@@ -0,0 +1,62 @@
+import torch
+import torch.nn as nn
+
+def supported_hyperparameters():
+    return {'lr', 'momentum'}
+
+class Net(nn.Module):
+    def __init__(self, in_shape: tuple, out_shape: tuple, prm: dict, device: torch.device) -> None:
+        super().__init__()
+        self.device = device
+        self.in_channels = in_shape[1]
+        self.num_classes = out_shape[0]
+        self.learning_rate = prm.get('lr', 0.01)
+        self.momentum = prm.get('momentum', 0.9)
+
+        self.features = nn.Sequential(
+            nn.Conv2d(self.in_channels, 64, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+        )
+
+        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
+
+        self.classifier = nn.Sequential(
+            nn.Dropout(p=prm.get('dropout', 0.5)),
+            nn.Linear(128 * 6 * 6, 4096),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=prm.get('dropout', 0.5)),
+            nn.Linear(4096, self.num_classes),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.features(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.classifier(x)
+        return x
+
+    def train_setup(self, prm):
+        self.to(self.device)
+        self.criteria = (nn.CrossEntropyLoss().to(self.device),)
+        self.optimizer = torch.optim.SGD(
+            self.parameters(),
+            lr=prm.get('lr', 0.01),
+            momentum=prm.get('momentum', 0.9),
+        )
+
+    def learn(self, train_data):
+        self.train()
+        for inputs, labels in train_data:
+            inputs = inputs.to(self.device)
+            labels = labels.to(self.device)
+            self.optimizer.zero_grad()
+            outputs = self(inputs)
+            loss = self.criteria[0](outputs, labels)
+            loss.backward()
+            self.optimizer.step()