神经架构搜索 NAS：自动设计网络

import torch
import torch.nn as nn

class NASCell(nn.Module):
    """NAS 搜索出的基本计算单元（cell）
    每个 cell 由多个节点组成，每个节点是对前驱节点的
    某种运算（卷积、池化、跳跃连接等）的组合"""
    def __init__(self, in_channels, out_channels, steps=4):
        super().__init__()
        self.steps = steps
        self.preprocessors = nn.ModuleList([
            nn.Conv2d(in_channels, out_channels, 1) for _ in range(2)
        ])
        # 每个节点有 2 个输入边
        self.edges = nn.ModuleList([
            nn.ModuleDict({
                f"op_{j}": self._make_mix_op(out_channels)
                for j in range(2)
            })
            for _ in range(steps)
        ])
    
    def _make_mix_op(self, channels):
        """混合运算：包含多种候选操作"""
        return nn.ModuleDict({
            "conv_3x3": nn.Conv2d(channels, channels, 3, padding=1),
            "conv_5x5": nn.Conv2d(channels, channels, 5, padding=2),
            "max_pool": nn.MaxPool2d(3, stride=1, padding=1),
            "avg_pool": nn.AvgPool2d(3, stride=1, padding=1),
            "skip_connect": nn.Identity(),
            "zero": nn.Conv2d(channels, channels, 1),
        })
    
    def forward(self, s0, s1):
        """s0, s1: 前两个节点的输出"""
        states = [self.preprocessors[0](s0), self.preprocessors[1](s1)]
        for i, edge_ops in enumerate(self.edges):
            # 每个节点选择两个前驱节点进行操作
            s_cur = edge_ops["op_0"](states[-2]) + edge_ops["op_1"](states[-1])
            states.append(s_cur)
        return torch.cat(states[-self.steps:], dim=1)

# NAS 双层优化的简化数学表述
import numpy as np

def bilevel_nas_formulation():
    """
    NAS 的双层优化问题：
    
    外层（架构搜索）:
        min_α  L_val(w*(α), α)
        s.t. w*(α) = argmin_w L_train(w, α)
    
    其中:
        α = 架构参数（离散选择或连续松弛）
        w = 网络权重
        L_train = 训练集损失
        L_val = 验证集损失
    
    这个公式揭示了一个关键洞察：
    架构的好坏取决于它训练后的验证表现，
    而训练效果又取决于架构本身。
    """
    print("NAS 双层优化:")
    print("  外层:  min_α L_val(w*(α), α)")
    print("  内层:  w*(α) = argmin_w L_train(w, α)")
    print()
    print("直接求解的计算量:")
    print("  每个候选架构 α 需要完整训练 w")
    print("  搜索 10^6 个架构 × 每个训练 100 epoch")
    print("  ≈ 10^8 GPU 小时（不可行）")
    print()
    print("加速策略:")
    print("  1. 权重共享（Weight Sharing）")
    print("  2. One-Shot / Supernet 方法")
    print("  3. 可微松弛（DARTS）")
    print("  4. 代理指标（Proxy Tasks）")
    print("  5. 预测器（Predictor-based）")

bilevel_nas_formulation()

# Cell-based 搜索空间的编码方式
# 用邻接矩阵 + 操作列表表示一个 DAG 结构的 cell
import torch

class CellEncoding:
    """Cell 的图编码：邻接矩阵 + 操作编码"""
    def __init__(self, num_nodes=4, num_ops=5):
        self.num_nodes = num_nodes
        self.num_ops = num_ops
        # 操作映射
        self.op_names = [
            "conv_3x3", "conv_5x5", "max_pool_3x3",
            "avg_pool_3x3", "skip_connect"
        ]
    
    def random_cell(self):
        """随机生成一个合法的 cell 编码"""
        # 邻接矩阵：上三角矩阵（DAG 约束）
        adj = torch.zeros(self.num_nodes + 2, self.num_nodes + 2)
        for i in range(self.num_nodes + 1):
            for j in range(i + 1, self.num_nodes + 2):
                # 每对节点之间有边（0/1）
                adj[i, j] = torch.bernoulli(torch.tensor(0.5))
        
        # 每个有边的操作类型
        ops = torch.randint(0, self.num_ops, (self.num_nodes + 2,))
        ops[0] = 0  # 输入节点
        ops[-1] = 0  # 输出节点
        
        return {"adj": adj, "ops": ops}
    
    def decode_to_graph(self, encoding):
        """将编码转换为可执行计算图"""
        adj = encoding["adj"]
        ops = encoding["ops"]
        
        nodes = []
        # 输入节点
        for i in range(2, self.num_nodes + 2):
            predecessors = torch.where(adj[:i, i] == 1)[0]
            node_ops = [self.op_names[ops[p].item()] for p in predecessors]
            nodes.append({
                "predecessors": predecessors.tolist(),
                "operations": node_ops
            })
        return nodes

# 示例：生成并解析一个随机 cell
encoder = CellEncoding(num_nodes=4, num_ops=5)
cell = encoder.random_cell()
graph = encoder.decode_to_graph(cell)
for i, node in enumerate(graph):
    print(f"Node {i}: inputs={node['predecessors']}, ops={node['operations']}")

# 不同搜索空间的复杂度对比
import math

def analyze_search_space():
    """量化分析不同搜索空间的大小"""
    configs = {
        "链式 (8 层, 5 种操作)": {
            "layers": 8,
            "choices_per_layer": 5,
            "formula": "5^8"
        },
        "Cell (4 节点, 5 操作, 2 输入/节点)": {
            "intermediate_nodes": 4,
            "choices_per_node": 5,
            "input_pairs": math.comb(4 + 1, 2),  # C(5,2)
            "formula": "10^(5^4)"
        },
        "层次化 (3 层, 每层 2 选择)": {
            "levels": 3,
            "choices_per_level": 2,
            "formula": "2^(2^3)"
        },
        "宏-微观联合 (20 个 cell, 每 cell 100 种)": {
            "macro": "stack pattern",
            "micro": "cell structure",
            "formula": "~10^20"
        },
    }
    
    print("搜索空间大小对比:")
    print("=" * 60)
    for name, cfg in configs.items():
        size = eval(cfg["formula"])
        print(f"  {name}")
        print(f"    大小: {size:.2e}")
        print(f"    公式: {cfg['formula']}")
        print(f"    log2(空间): {math.log2(size):.1f} bits")
        print()
    
    print("结论:")
    print("  - 链式空间小但表达能力有限")
    print("  - Cell-based 是性价比最高的选择")
    print("  - 联合搜索需要强力评估策略")

analyze_search_space()

# 强化学习搜索策略：基于 PPO 的架构控制器
import torch
import torch.nn as nn
import torch.nn.functional as F

class RLController(nn.Module):
    """RNN 控制器：逐步生成架构描述"""
    def __init__(self, num_layers, choices_per_layer, hidden_size=64):
        super().__init__()
        self.num_layers = num_layers
        self.choices_per_layer = choices_per_layer
        self.hidden_size = hidden_size
        
        self.embedding = nn.Embedding(choices_per_layer, hidden_size)
        self.lstm = nn.LSTMCell(hidden_size, hidden_size)
        self.predictor = nn.Linear(hidden_size, choices_per_layer)
        self.temperature = nn.Parameter(torch.tensor(5.0))
    
    def sample_architecture(self, batch_size=1):
        """采样一批架构"""
        inputs = torch.zeros(batch_size, self.hidden_size)
        h, c = torch.zeros(batch_size, self.hidden_size), torch.zeros(batch_size, self.hidden_size)
        
        actions = []
        log_probs = []
        entropies = []
        
        for i in range(self.num_layers):
            h, c = self.lstm(inputs, (h, c))
            logits = self.predictor(h) / self.temperature
            probs = F.softmax(logits, dim=-1)
            
            # Gumbel-Softmax 或直接采样
            action = torch.multinomial(probs, 1).squeeze(-1)
            lp = F.log_softmax(logits, dim=-1).gather(1, action.unsqueeze(1)).squeeze()
            entropy = -(probs * F.log_softmax(logits, dim=-1)).sum(dim=-1)
            
            actions.append(action)
            log_probs.append(lp)
            entropies.append(entropy)
            
            inputs = self.embedding(action)
        
        return torch.stack(actions, dim=1), torch.stack(log_probs, dim=1), torch.stack(entropies, dim=1)

# 进化算法搜索策略
import random
from copy import deepcopy

class EvolutionaryNAS:
    """简化版正则进化 NAS"""
    def __init__(self, population_size=50, tournament_size=10, mutation_rate=0.3):
        self.population_size = population_size
        self.tournament_size = tournament_size
        self.mutation_rate = mutation_rate
        self.population = []  # [(arch_encoding, accuracy, age), ...]
    
    def mutate(self, arch):
        """对架构进行随机变异"""
        mutated = deepcopy(arch)
        # 变异类型 1: 修改某个节点的操作
        if random.random() < self.mutation_rate:
            node_idx = random.randint(0, len(mutated["ops"]) - 1)
            mutated["ops"][node_idx] = random.randint(0, 4)
        # 变异类型 2: 添加/删除边
        if random.random() < self.mutation_rate:
            i = random.randint(0, mutated["adj"].shape[0] - 2)
            j = random.randint(i + 1, mutated["adj"].shape[0] - 1)
            mutated["adj"][i, j] = 1 - mutated["adj"][i, j]
        return mutated
    
    def tournament_selection(self):
        """锦标赛选择：随机选 tournament_size 个，返回最优的"""
        candidates = random.sample(self.population, self.tournament_size)
        return max(candidates, key=lambda x: x[1])  # 按 accuracy 选择
    
    def step(self, evaluate_fn):
        """执行一代进化"""
        # 选择父代
        parent = self.tournament_selection()
        child_arch = self.mutate(parent[0])
        child_acc = evaluate_fn(child_arch)
        
        # 正则化进化：替换最老的个体
        oldest = min(self.population, key=lambda x: x[2])
        idx = self.population.index(oldest)
        self.population[idx] = (child_arch, child_acc, 0)
        
        # 所有个体年龄 +1
        for i in range(len(self.population)):
            self.population[i] = (self.population[i][0], self.population[i][1], 
                                  self.population[i][2] + 1)
        
        best = max(self.population, key=lambda x: x[1])
        return best
    
    def search(self, evaluate_fn, generations=200):
        """初始化种群并开始搜索"""
        for _ in range(self.population_size):
            arch = {"ops": [random.randint(0, 4) for _ in range(6)],
                    "adj": torch.bernoulli(torch.full((6, 6), 0.3)).triu(1)}
            acc = evaluate_fn(arch)
            self.population.append((arch, acc, 0))
        
        best_acc = 0
        for gen in range(generations):
            best = self.step(evaluate_fn)
            if best[1] > best_acc:
                best_acc = best[1]
                print(f"Gen {gen}: new best accuracy = {best_acc:.4f}")

# 模拟评估函数（实际应训练并验证）
def mock_evaluate(arch):
    return 0.85 + 0.1 * random.random()  # 模拟 85%-95% 的准确率

# One-Shot Supernet 实现
import torch
import torch.nn as nn
import torch.nn.functional as F

class OneShotSupernet(nn.Module):
    """One-Shot NAS 超网络：所有操作同时训练"""
    def __init__(self, num_classes=10, num_stages=3, channels=64):
        super().__init__()
        self.num_stages = num_stages
        
        # 每个 stage 有多个候选操作
        self.stem = nn.Sequential(
            nn.Conv2d(3, channels, 3, padding=1),
            nn.BatchNorm2d(channels),
            nn.ReLU()
        )
        
        self.stages = nn.ModuleList()
        for s in range(num_stages):
            stage = OneShotStage(channels, s > 0)
            self.stages.append(stage)
            if s > 0:
                channels *= 2  # 下采样后通道翻倍
        
        self.global_pool = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Linear(channels, num_classes)
    
    def forward(self, x, sample_path=None):
        """
        sample_path: 指定采样路径 (训练超网络时为 None，评估子网络时为具体路径)
        """
        x = self.stem(x)
        for stage in self.stages:
            x = stage(x, sample_path)
        x = self.global_pool(x).flatten(1)
        return self.classifier(x)


class OneShotStage(nn.Module):
    """一个 stage 包含所有候选操作"""
    def __init__(self, channels, downsample=False):
        super().__init__()
        stride = 2 if downsample else 1
        self.ops = nn.ModuleDict({
            "conv3x3": nn.Conv2d(channels, channels * (2 if downsample else 1), 
                                  3, stride=stride, padding=1),
            "conv5x5": nn.Conv2d(channels, channels * (2 if downsample else 1), 
                                  5, stride=stride, padding=2),
            "maxpool": nn.MaxPool2d(3, stride=stride, padding=1),
            "avgpool": nn.AvgPool2d(3, stride=stride, padding=1),
        })
    
    def forward(self, x, sample_path=None):
        if sample_path is not None:
            # 评估模式：只运行指定操作
            return self.ops[sample_path](x)
        else:
            # 训练模式：所有操作共享梯度
            return sum(self.ops[name](x) for name in self.ops) / len(self.ops)

# 超网络训练 + 子架构采样评估
import random
import torch.optim as optim

def train_supernet(supernet, dataloader, epochs=50):
    """训练超网络：每次随机采样一个子架构路径"""
    optimizer = optim.Adam(supernet.parameters(), lr=0.025, weight_decay=3e-4)
    
    for epoch in range(epochs):
        for batch_x, batch_y in dataloader:
            optimizer.zero_grad()
            
            # 随机采样路径
            op_names = list(supernet.stages[0].ops.keys())
            sample_path = [random.choice(op_names) for _ in range(supernet.num_stages)]
            
            # 前向传播
            logits = supernet(batch_x, sample_path=sample_path)
            loss = F.cross_entropy(logits, batch_y)
            
            # 反向传播：梯度只流向当前采样路径涉及的操作
            loss.backward()
            optimizer.step()
        
        if epoch % 10 == 0:
            print(f"Epoch {epoch}: loss = {loss.item():.4f}")


def evaluate_architecture(supernet, arch_encoding, val_loader):
    """评估特定架构在验证集上的表现"""
    supernet.eval()
    correct = 0
    total = 0
    
    with torch.no_grad():
        for x, y in val_loader:
            logits = supernet(x, sample_path=arch_encoding)
            preds = logits.argmax(dim=1)
            correct += (preds == y).sum().item()
            total += y.size(0)
    
    return correct / total


def search(supernet, val_loader, num_samples=100):
    """从训练好的超网络中搜索最佳架构"""
    op_names = list(supernet.stages[0].ops.keys())
    best_acc = 0
    best_arch = None
    
    for _ in range(num_samples):
        arch = [random.choice(op_names) for _ in range(supernet.num_stages)]
        acc = evaluate_architecture(supernet, arch, val_loader)
        if acc > best_acc:
            best_acc = acc
            best_arch = arch
            print(f"New best: {best_arch} -> {best_acc:.4f}")
    
    return best_arch, best_acc

import torch
import torch.nn as nn
import torch.nn.functional as F

class MixedOp(nn.Module):
    """混合操作：所有候选操作的加权组合"""
    def __init__(self, channels, stride=1):
        super().__init__()
        self._ops = nn.ModuleList([
            nn.Conv2d(channels, channels, 1, stride=stride),
            nn.Conv2d(channels, channels, 3, stride=stride, padding=1),
            nn.Conv2d(channels, channels, 5, stride=stride, padding=2),
            nn.MaxPool2d(3, stride=stride, padding=1),
            nn.AvgPool2d(3, stride=stride, padding=1),
            nn.Identity() if stride == 1 else nn.Conv2d(channels, channels, 1, stride=stride),
        ])
    
    def forward(self, x, weights):
        """
        weights: 长度为 num_ops 的向量，经 softmax 归一化
        """
        return sum(w * op(x) for w, op in zip(weights, self._ops))


class DARTSCell(nn.Module):
    """DARTS 可微 Cell"""
    def __init__(self, steps=4, multiplier=4, channels=64):
        super().__init__()
        self.steps = steps
        self.multiplier = multiplier
        self._k = sum(1 for i in range(steps) for _ in range(i + 2))
        
        # 架构参数（可学习）
        self.alphas_normal = nn.Parameter(1e-3 * torch.randn(self._k, 6))
        self.alphas_reduce = nn.Parameter(1e-3 * torch.randn(self._k, 6))
        
        # 混合操作
        self._ops = nn.ModuleList()
        for i in range(self._k):
            stride = 2 if i < 2 else 1  # 前两个是 reduction 边
            self._ops.append(MixedOp(channels, stride))
    
    def forward(self, s0, s1, normal_weights, reduce_weights):
        """
        s0, s1: 前两个节点的输出
        normal_weights, reduce_weights: 归一化的架构权重
        """
        states = [s0, s1]
        offset = 0
        for i in range(self.steps):
            # 判断是 normal 还是 reduce 边
            w = reduce_weights[offset:offset + i + 2] if i < 2 else normal_weights[offset:offset + i + 2]
            s = sum(self._ops[offset + j](states[j], w[j]) for j in range(i + 2))
            offset += i + 2
            states.append(s)
        return torch.cat(states[-self.steps:], dim=1)

class DARTS(nn.Module):
    """完整 DARTS 模型：可微架构搜索"""
    def __init__(self, C=16, num_classes=10, layers=8, steps=4):
        super().__init__()
        self.layers = layers
        self.stem = nn.Sequential(
            nn.Conv2d(3, C, 3, padding=1, bias=False),
            nn.BatchNorm2d(C)
        )
        
        # 预定义 reduction 位置（1/3 和 2/3 处）
        self.cells = nn.ModuleList()
        C_prev_prev, C_prev, C_curr = C, C, C
        
        for i in range(layers):
            if i in [layers // 3, 2 * layers // 3]:
                C_curr *= 2
            
            cell = DARTSCell(steps, 4, C_curr)
            self.cells.append(cell)
            C_prev = C_curr * steps
        
        self.global_pool = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Linear(C_prev, num_classes)
    
    def normal_weights(self):
        return F.softmax(self.cells[0].alphas_normal, dim=-1)
    
    def reduce_weights(self):
        return F.softmax(self.cells[0].alphas_reduce, dim=-1)
    
    def forward(self, x):
        s0 = s1 = self.stem(x)
        for i, cell in enumerate(self.cells):
            s0, s1 = s1, cell(s0, s1, self.normal_weights(), self.reduce_weights())
        out = self.global_pool(s1).flatten(1)
        return self.classifier(out)


# DARTS 交替优化训练循环
def train_darts_step(model, train_loader, val_loader, opt_w, opt_alpha, device):
    """交替优化：先优化权重 w，再优化架构参数 α"""
    # Step 1: 优化权重 w (在训练集上)
    model.train()
    for batch_x, batch_y in train_loader:
        batch_x, batch_y = batch_x.to(device), batch_y.to(device)
        logits = model(batch_x)
        loss_w = F.cross_entropy(logits, batch_y)
        opt_w.zero_grad()
        loss_w.backward()
        opt_w.step()
    
    # Step 2: 优化架构参数 α (在验证集上)
    for batch_x, batch_y in val_loader:
        batch_x, batch_y = batch_x.to(device), batch_y.to(device)
        logits = model(batch_x)
        loss_alpha = F.cross_entropy(logits, batch_y)
        opt_alpha.zero_grad()
        loss_alpha.backward()
        opt_alpha.step()

# MNASNet 风格的延迟感知奖励函数
import time
import torch

class LatencyAwareReward:
    """MNASNet 的延迟感知奖励计算"""
    def __init__(self, target_latency_ms=78.0, weight_w=-0.07):
        self.target_latency_ms = target_latency_ms
        self.w = weight_w
    
    def measure_latency(self, model, input_shape=(1, 3, 224, 224), warmup=10, trials=100):
        """在目标设备上测量推理延迟"""
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = model.to(device).eval()
        x = torch.randn(input_shape).to(device)
        
        # Warmup
        with torch.no_grad():
            for _ in range(warmup):
                _ = model(x)
        
        # Timing
        torch.cuda.synchronize() if torch.cuda.is_available() else None
        start = time.perf_counter()
        with torch.no_grad():
            for _ in range(trials):
                _ = model(x)
        torch.cuda.synchronize() if torch.cuda.is_available() else None
        end = time.perf_counter()
        
        return (end - start) / trials * 1000  # ms
    
    def compute_reward(self, accuracy, latency_ms):
        """MNASNet 奖励函数: R = ACC * (LAT / TARGET)^w"""
        if latency_ms <= 0:
            return accuracy
        ratio = latency_ms / self.target_latency_ms
        reward = accuracy * (ratio ** self.w)
        return reward
    
    def __call__(self, model, accuracy, input_shape=(1, 3, 224, 224)):
        latency = self.measure_latency(model, input_shape)
        reward = self.compute_reward(accuracy, latency)
        print(f"Accuracy: {accuracy:.4f}, Latency: {latency:.1f}ms, Reward: {reward:.4f}")
        return reward, latency

# EfficientNet 复合缩放实现
import torch
import torch.nn as nn
import math

def efficientnet_coefficients():
    """
    EfficientNet 复合缩放参数
    通过小规模网格搜索确定最优 (α, β, γ)
    """
    # B0 到 B7 的缩放系数
    configs = {
        "b0": {"width": 1.0, "depth": 1.0, "resolution": 224, "dropout": 0.2},
        "b1": {"width": 1.0, "depth": 1.1, "resolution": 240, "dropout": 0.2},
        "b2": {"width": 1.1, "depth": 1.2, "resolution": 260, "dropout": 0.3},
        "b3": {"width": 1.2, "depth": 1.4, "resolution": 300, "dropout": 0.3},
        "b4": {"width": 1.4, "depth": 1.8, "resolution": 380, "dropout": 0.4},
        "b5": {"width": 1.6, "depth": 2.2, "resolution": 456, "dropout": 0.4},
        "b6": {"width": 1.8, "depth": 2.6, "resolution": 528, "dropout": 0.5},
        "b7": {"width": 2.0, "depth": 3.1, "resolution": 600, "dropout": 0.5},
    }
    return configs


class EfficientNetBlock(nn.Module):
    """MBConv 倒置瓶颈块 (Mobile Inverted Bottleneck)"""
    def __init__(self, in_channels, out_channels, kernel_size, 
                 expand_ratio=1, stride=1, se_ratio=0.25):
        super().__init__()
        mid_channels = in_channels * expand_ratio
        self.use_residual = stride == 1 and in_channels == out_channels
        
        layers = []
        # 1x1 扩展卷积 (扩展通道)
        if expand_ratio != 1:
            layers.extend([
                nn.Conv2d(in_channels, mid_channels, 1, bias=False),
                nn.BatchNorm2d(mid_channels),
                nn.SiLU()
            ])
        
        # 深度可分离卷积 (KxK)
        layers.extend([
            nn.Conv2d(mid_channels, mid_channels, kernel_size, 
                     stride=stride, padding=kernel_size//2,
                     groups=mid_channels, bias=False),
            nn.BatchNorm2d(mid_channels),
            nn.SiLU()
        ])
        
        # Squeeze-and-Excitation (注意力)
        se_channels = max(1, int(in_channels * se_ratio))
        layers.extend([
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(mid_channels, se_channels, 1),
            nn.SiLU(),
            nn.Conv2d(se_channels, mid_channels, 1),
            nn.Sigmoid()
        ])
        
        # 1x1 投影卷积 (压缩通道)
        layers.extend([
            nn.Conv2d(mid_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels)
        ])
        
        self.block = nn.Sequential(*layers)
    
    def forward(self, x):
        out = self.block(x)
        # 分离 Squeeze-and-Excitation 的输出
        return out + x if self.use_residual else out


class CompoundScaler:
    """EfficientNet 复合缩放"""
    def __init__(self, phi=1.0, alpha=1.2, beta=1.1, gamma=1.15):
        self.phi = phi
        self.alpha = alpha  # 深度系数
        self.beta = beta    # 宽度系数
        self.gamma = gamma  # 分辨率系数
    
    def scale(self, depth, width, resolution):
        """应用复合缩放"""
        new_depth = int(math.ceil(self.alpha ** self.phi * depth))
        new_width = int(math.ceil(self.beta ** self.phi * width) / 8) * 8
        new_resolution = int(math.ceil(self.gamma ** self.phi * resolution) / 32) * 32
        return new_depth, new_width, new_resolution

# 示例：从 B0 缩放到 B3
scaler = CompoundScaler(phi=3.0, alpha=1.2, beta=1.1, gamma=1.15)
for i in range(4):
    s = CompoundScaler(phi=float(i))
    d, w, r = s.scale(1, 32, 224)
    print(f"B{i}: depth_mult={s.alpha**i:.2f}, width={w}, resolution={r}")

# search_space.json: NNI 搜索空间定义
"""
{
    "conv1_kernel": {"_type": "choice", "_value": [3, 5, 7]},
    "conv1_channels": {"_type": "choice", "_value": [16, 32, 64]},
    "conv2_kernel": {"_type": "choice", "_value": [3, 5, 7]},
    "conv2_channels": {"_type": "choice", "_value": [32, 64, 128]},
    "conv3_kernel": {"_type": "choice", "_value": [3, 5]},
    "conv3_channels": {"_type": "choice", "_value": [64, 128, 256]},
    "dropout_rate": {"_type": "uniform", "_value": [0.1, 0.5]},
    "use_bn": {"_type": "choice", "_value": [true, false]},
    "fc_hidden": {"_type": "choice", "_value": [64, 128, 256, 512]}
}

这个搜索空间定义了 3*3*3*3*2*3 * 连续 * 2 * 4 ≈ 3888 种离散组合
加上连续的 dropout_rate，总空间远大于此。
TPE tuner 可以在 50-100 次 trial 内找到接近最优的配置。
"""
import json

search_space = {
    "conv1_kernel": {"_type": "choice", "_value": [3, 5, 7]},
    "conv1_channels": {"_type": "choice", "_value": [16, 32, 64]},
    "conv2_kernel": {"_type": "choice", "_value": [3, 5, 7]},
    "conv2_channels": {"_type": "choice", "_value": [32, 64, 128]},
    "conv3_kernel": {"_type": "choice", "_value": [3, 5]},
    "conv3_channels": {"_type": "choice", "_value": [64, 128, 256]},
    "dropout_rate": {"_type": "uniform", "_value": [0.1, 0.5]},
    "use_bn": {"_type": "choice", "_value": [True, False]},
    "fc_hidden": {"_type": "choice", "_value": [64, 128, 256, 512]},
}

with open("search_space.json", "w") as f:
    json.dump(search_space, f, indent=2)
print("search_space.json 已生成")
print(f"搜索空间大小估算: {3*3*3*3*2*3*2*4} 种离散组合")

# model.py: NNI 可搜索模型 + 训练脚本
import nni
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim


class SearchableCNN(nn.Module):
    """使用 NNI 参数的可搜索 CNN"""
    def __init__(self, params):
        super().__init__()
        self.conv1 = nn.Conv2d(3, params["conv1_channels"], 
                                params["conv1_kernel"], padding=params["conv1_kernel"]//2)
        self.conv2 = nn.Conv2d(params["conv1_channels"], params["conv2_channels"],
                                params["conv2_kernel"], padding=params["conv2_kernel"]//2)
        self.conv3 = nn.Conv2d(params["conv2_channels"], params["conv3_channels"],
                                params["conv3_kernel"], padding=params["conv3_kernel"]//2)
        
        if params["use_bn"]:
            self.bn1 = nn.BatchNorm2d(params["conv1_channels"])
            self.bn2 = nn.BatchNorm2d(params["conv2_channels"])
            self.bn3 = nn.BatchNorm2d(params["conv3_channels"])
        else:
            self.bn1 = self.bn2 = self.bn3 = nn.Identity()
        
        self.pool = nn.MaxPool2d(2, 2)
        self.dropout = nn.Dropout(params["dropout_rate"])
        
        # 自适应计算 FC 输入维度
        self.fc1 = nn.Linear(params["conv3_channels"] * 4 * 4, params["fc_hidden"])
        self.fc2 = nn.Linear(params["fc_hidden"], 10)
    
    def forward(self, x):
        x = self.pool(F.relu(self.bn1(self.conv1(x))))
        x = self.pool(F.relu(self.bn2(self.conv2(x))))
        x = self.pool(F.relu(self.bn3(self.conv3(x))))
        x = x.view(x.size(0), -1)
        x = self.dropout(x)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x


def train_and_evaluate(params, epochs=10):
    """训练模型并返回验证准确率"""
    model = SearchableCNN(params)
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()
    
    # 模拟训练循环（实际应使用真实数据加载器）
    for epoch in range(epochs):
        model.train()
        # 模拟训练步
        dummy_x = torch.randn(32, 3, 32, 32)
        dummy_y = torch.randint(0, 10, (32,))
        logits = model(dummy_x)
        loss = criterion(logits, dummy_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    # 评估
    model.eval()
    with torch.no_grad():
        test_x = torch.randn(1000, 3, 32, 32)
        test_y = torch.randint(0, 10, (1000,))
        logits = model(test_x)
        acc = (logits.argmax(1) == test_y).float().mean().item()
    
    return acc


def main():
    """NNI Trial 入口"""
    # 获取 NNI 传入的超参数
    params = nni.get_next_parameter()
    
    # 训练并评估
    accuracy = train_and_evaluate(params)
    
    # 报告结果给 NNI
    nni.report_final_result(accuracy)
    print(f"Trial 完成: accuracy = {accuracy:.4f}")
    print(f"Params: {params}")


if __name__ == "__main__":
    main()