21 - 元学习 (Meta-Learning)¶

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

class MAML(nn.Module):  # 继承nn.Module定义神经网络层
    def __init__(self, model, inner_lr=0.01, num_inner_steps=1):  # __init__构造方法，创建对象时自动调用
        super().__init__()  # super()调用父类方法
        self.model = model
        self.inner_lr = inner_lr
        self.num_inner_steps = num_inner_steps

    def inner_loop(self, support_x, support_y):
        """内循环：在支持集上适应"""
        # 创建任务特定的参数副本
        fast_weights = [p.clone() for p in self.model.parameters()]

        for _ in range(self.num_inner_steps):
            # 前向传播
            logits = self.forward_with_params(support_x, fast_weights)
            loss = F.cross_entropy(logits, support_y)

            # 计算梯度
            grads = torch.autograd.grad(loss, fast_weights, create_graph=True)

            # 更新快速权重
            fast_weights = [w - self.inner_lr * g
                          for w, g in zip(fast_weights, grads)]  # zip按位置配对多个可迭代对象

        return fast_weights

    def forward_with_params(self, x, params):
        """使用指定参数前向传播（简化为2层MLP示例）

        注意：这里使用简单的MLP结构来演示MAML核心原理。
        若使用SimpleCNN等复杂模型，需通过torch.nn.utils中的
        功能性API遍历每一层进行参数替换。
        """
        x = x.view(x.size(0), -1)  # 展平输入
        x = F.linear(x, params[0], params[1])
        x = F.relu(x)
        x = F.linear(x, params[2], params[3])
        return x

    def forward(self, batch_tasks, meta_optimizer):
        """元训练"""
        meta_loss = 0

        for task in batch_tasks:
            support_x, support_y = task['support']
            query_x, query_y = task['query']

            # 内循环适应
            fast_weights = self.inner_loop(support_x, support_y)

            # 在查询集上评估
            query_logits = self.forward_with_params(query_x, fast_weights)
            task_loss = F.cross_entropy(query_logits, query_y)
            meta_loss += task_loss

        # 元更新
        meta_optimizer.zero_grad()
        meta_loss.backward()  # 反向传播计算梯度
        meta_optimizer.step()  # 根据梯度更新模型参数

        return meta_loss.item()  # .item()将单元素张量转为Python数值

# 简单的4层CNN
class SimpleCNN(nn.Module):
    def __init__(self, in_channels=3, num_classes=5):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(in_channels, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Linear(64, num_classes)

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return self.classifier(x)

class PrototypicalNetwork(nn.Module):
    def __init__(self, encoder):
        super().__init__()
        self.encoder = encoder

    def euclidean_distance(self, x, y):
        """计算欧氏距离"""
        n = x.size(0)
        m = y.size(0)
        d = x.size(1)

        x = x.unsqueeze(1).expand(n, m, d)  # 链式调用，连续执行多个方法  # unsqueeze增加一个维度
        y = y.unsqueeze(0).expand(n, m, d)

        return torch.pow(x - y, 2).sum(2)

    def forward(self, support_x, support_y, query_x, n_way, k_shot):
        """
        support_x: (n_way * k_shot, ...)
        support_y: (n_way * k_shot,)
        query_x: (n_query, ...)
        """
        # 编码
        support_embed = self.encoder(support_x)  # (n_way*k_shot, embed_dim)
        query_embed = self.encoder(query_x)      # (n_query, embed_dim)

        # 计算原型
        prototypes = []
        for k in range(n_way):
            class_mask = (support_y == k)
            class_embed = support_embed[class_mask]
            prototype = class_embed.mean(dim=0)
            prototypes.append(prototype)

        prototypes = torch.stack(prototypes)  # (n_way, embed_dim)

        # 计算查询样本到原型的距离
        distances = self.euclidean_distance(query_embed, prototypes)

        # 转换为概率（距离越小，概率越大）
        logits = -distances

        return logits

# 训练
proto_net = PrototypicalNetwork(SimpleCNN())
optimizer = torch.optim.Adam(proto_net.parameters(), lr=1e-3)

for episode in range(num_episodes):
    # 采样一个episode
    support_x, support_y, query_x, query_y = sample_episode(dataset, n_way, k_shot, n_query)

    # 前向传播
    logits = proto_net(support_x, support_y, query_x, n_way, k_shot)

    # 计算损失
    loss = F.cross_entropy(logits, query_y)

    # 反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

class EpisodeSampler:
    def __init__(self, dataset, n_way, k_shot, n_query):
        self.dataset = dataset
        self.n_way = n_way
        self.k_shot = k_shot
        self.n_query = n_query

        # 按类别组织数据
        self.class_to_indices = {}
        for idx, label in enumerate(dataset.labels):  # enumerate同时获取索引和元素
            if label not in self.class_to_indices:
                self.class_to_indices[label] = []
            self.class_to_indices[label].append(idx)

        self.classes = list(self.class_to_indices.keys())

    def sample_episode(self):
        """采样一个episode"""
        # 随机选择n_way个类别
        episode_classes = np.random.choice(self.classes, self.n_way, replace=False)

        support_x, support_y = [], []
        query_x, query_y = [], []

        for new_label, old_label in enumerate(episode_classes):
            # 获取该类所有样本索引
            indices = self.class_to_indices[old_label]

            # 随机选择k_shot + n_query个样本
            selected = np.random.choice(indices, self.k_shot + self.n_query, replace=False)

            # 分为支持集和查询集
            support_indices = selected[:self.k_shot]
            query_indices = selected[self.k_shot:]

            # 收集样本
            for idx in support_indices:
                support_x.append(self.dataset[idx][0])
                support_y.append(new_label)

            for idx in query_indices:
                query_x.append(self.dataset[idx][0])
                query_y.append(new_label)

        return (torch.stack(support_x), torch.tensor(support_y),
                torch.stack(query_x), torch.tensor(query_y))