04-神经架构搜索（NAS）¶

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

# 候选操作集合
OPS = {
    'sep_conv_3x3': lambda C: SepConv(C, C, 3, 1),  # lambda匿名函数，简洁定义单行函数
    'sep_conv_5x5': lambda C: SepConv(C, C, 5, 2),
    'max_pool_3x3': lambda C: nn.MaxPool2d(3, 1, 1),
    'avg_pool_3x3': lambda C: nn.AvgPool2d(3, 1, 1),
    'skip_connect': lambda C: nn.Identity(),
    'none':         lambda C: Zero(),
}

class SepConv(nn.Module):  # 继承nn.Module定义神经网络层
    """深度可分离卷积"""
    def __init__(self, C_in, C_out, kernel, padding):  # __init__构造方法，创建对象时自动调用
        super().__init__()  # super()调用父类方法
        self.op = nn.Sequential(
            nn.Conv2d(C_in, C_in, kernel, 1, padding, groups=C_in, bias=False),
            nn.Conv2d(C_in, C_out, 1, bias=False),
            nn.BatchNorm2d(C_out),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        return self.op(x)

class Zero(nn.Module):
    def forward(self, x):
        return x * 0

class MixedOp(nn.Module):
    """混合操作：所有候选操作的加权和"""
    def __init__(self, C):
        super().__init__()
        self.ops = nn.ModuleList([
            OPS[name](C) for name in OPS.keys()  # 列表推导式，简洁创建列表
        ])

    def forward(self, x, weights):
        return sum(w * op(x) for w, op in zip(weights, self.ops))  # zip按位置配对多个可迭代对象

class DARTSCell(nn.Module):
    """一个DARTS搜索单元"""
    def __init__(self, C, n_nodes=4):
        super().__init__()
        self.n_nodes = n_nodes
        self.ops = nn.ModuleList()
        # 每对节点之间都有一组混合操作
        for i in range(n_nodes):
            for j in range(i + 2):  # 包含2个输入节点
                self.ops.append(MixedOp(C))
        # 拼接n_nodes个中间节点后通道数为C*n_nodes，需要降回C
        self.channel_reduce = nn.Sequential(
            nn.Conv2d(C * n_nodes, C, 1, bias=False),
            nn.BatchNorm2d(C),
        )

    def forward(self, s0, s1, alphas):
        states = [s0, s1]
        offset = 0
        for i in range(self.n_nodes):
            s = sum(
                self.ops[offset + j](states[j], F.softmax(alphas[offset + j], dim=-1))
                for j in range(len(states))
            )
            offset += len(states)
            states.append(s)
        # 拼接中间节点后用1x1卷积将通道数从C*n_nodes降回C
        return self.channel_reduce(torch.cat(states[2:], dim=1))

class DARTSNetwork(nn.Module):
    """简化的DARTS搜索网络"""
    def __init__(self, C=16, n_classes=10, n_cells=4, n_nodes=4):
        super().__init__()
        self.stem = nn.Sequential(
            nn.Conv2d(3, C, 3, 1, 1, bias=False),
            nn.BatchNorm2d(C),
        )
        self.cells = nn.ModuleList([
            DARTSCell(C, n_nodes) for _ in range(n_cells)
        ])
        self.classifier = nn.Linear(C, n_classes)

        # 架构参数
        n_edges = sum(i + 2 for i in range(n_nodes))
        n_ops = len(OPS)
        self.alphas = nn.Parameter(torch.randn(n_edges, n_ops) * 1e-3)

    def forward(self, x):
        s0 = s1 = self.stem(x)
        for cell in self.cells:
            s0, s1 = s1, cell(s0, s1, self.alphas)
        out = F.adaptive_avg_pool2d(s1, 1).flatten(1)  # 链式调用，连续执行多个方法
        return self.classifier(out)

# 训练示意
def train_darts():
    model = DARTSNetwork(C=16, n_classes=10)
    w_optim = torch.optim.SGD(
        [p for n, p in model.named_parameters() if n != 'alphas'],
        lr=0.025, momentum=0.9, weight_decay=3e-4
    )
    a_optim = torch.optim.Adam([model.alphas], lr=3e-4, weight_decay=1e-3)

    # 交替优化（伪代码）
    for epoch in range(50):
        for x_train, y_train, x_val, y_val in zip_loaders():
            # 步骤1: 更新网络权重w
            w_optim.zero_grad()
            loss_train = F.cross_entropy(model(x_train), y_train)
            loss_train.backward()  # 反向传播计算梯度
            w_optim.step()

            # 步骤2: 更新架构参数α
            a_optim.zero_grad()
            loss_val = F.cross_entropy(model(x_val), y_val)
            loss_val.backward()
            a_optim.step()

    # 导出离散架构
    ops_names = list(OPS.keys())
    best_ops = [ops_names[i] for i in model.alphas.argmax(dim=-1)]
    print("搜索到的架构:", best_ops)