迁移学习（四）：小样本学习

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

for epoch in range(num_epochs):
    for episode in range(episodes_per_epoch):
        classes = sample(base_classes, N)
        support = sample_from_classes(classes, K)
        query   = sample_from_classes(classes, Q)

        prototypes = compute_prototypes(support)
        logits     = -distance(query, prototypes)
        loss       = cross_entropy(logits, query_labels)

        loss.backward()
        optimizer.step()

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

import torch
import torch.nn as nn
import numpy as np
from tqdm import tqdm

class ConvBlock(nn.Module):
    """Conv -> BN -> ReLU -> MaxPool，miniImageNet 的标配模块。"""
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
        )

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

class ProtoNetEncoder(nn.Module):
    """4 层 CNN 编码器，将 84x84 RGB 图像映射到 1600 维向量。"""
    def __init__(self, in_channels=3, hidden_dim=64):
        super().__init__()
        self.encoder = nn.Sequential(
            ConvBlock(in_channels, hidden_dim),   # 84 -> 42
            ConvBlock(hidden_dim, hidden_dim),    # 42 -> 21
            ConvBlock(hidden_dim, hidden_dim),    # 21 -> 10
            ConvBlock(hidden_dim, hidden_dim),    # 10 ->  5
        )

    def forward(self, x):
        return self.encoder(x).view(x.size(0), -1)

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

    def compute_prototypes(self, support_emb, support_labels, n_way):
        """对每个类别的支持嵌入取均值，生成原型。"""
        prototypes = [
            support_emb[support_labels == c].mean(dim=0)
            for c in range(n_way)
        ]
        return torch.stack(prototypes)  # (n_way, embed_dim)

    def forward(self, support_imgs, support_lbls, query_imgs, n_way):
        support_emb = self.encoder(support_imgs)
        query_emb   = self.encoder(query_imgs)
        prototypes  = self.compute_prototypes(support_emb, support_lbls, n_way)
        # 负欧氏距离作为 logits
        dists  = torch.cdist(query_emb, prototypes, p=2)
        return -dists

class EpisodeSampler:
    """从扁平化的 (data, labels) 中生成 N-way K-shot 的 episode。"""
    def __init__(self, data, labels, n_way, n_support, n_query, n_episodes):
        self.data, self.labels = data, labels
        self.n_way, self.n_support = n_way, n_support
        self.n_query, self.n_episodes = n_query, n_episodes
        self.classes = np.unique(labels)
        self.class_indices = {c: np.where(labels == c)[0] for c in self.classes}

    def __iter__(self):
        for _ in range(self.n_episodes):
            yield self._sample()

    def _sample(self):
        chosen = np.random.choice(self.classes, self.n_way, replace=False)
        s_imgs, s_lbls, q_imgs, q_lbls = [], [], [], []

        for new_label, c in enumerate(chosen):
            idxs = np.random.choice(
                self.class_indices[c],
                self.n_support + self.n_query,
                replace=False,
            )
            for idx in idxs[:self.n_support]:
                s_imgs.append(self.data[idx]); s_lbls.append(new_label)
            for idx in idxs[self.n_support:]:
                q_imgs.append(self.data[idx]); q_lbls.append(new_label)

        return (
            torch.stack([torch.FloatTensor(x) for x in s_imgs]),
            torch.LongTensor(s_lbls),
            torch.stack([torch.FloatTensor(x) for x in q_imgs]),
            torch.LongTensor(q_lbls),
        )

def train(model, train_data, train_lbls, val_data, val_lbls,
          n_way=5, n_support=5, n_query=15, n_episodes=100,
          num_epochs=50, lr=1e-3, device='cpu'):
    """Episode 训练循环，带周期性验证。"""
    model = model.to(device)
    optim = torch.optim.Adam(model.parameters(), lr=lr)
    crit = nn.CrossEntropyLoss()
    best = 0.0

    for epoch in range(num_epochs):
        # ---- 训练阶段 ----
        model.train()
        sampler = EpisodeSampler(train_data, train_lbls,
                                 n_way, n_support, n_query, n_episodes)
        loss_sum, acc_sum = 0.0, 0.0
        for s_img, s_lbl, q_img, q_lbl in tqdm(sampler, desc=f'Epoch {epoch+1}'):
            s_img, s_lbl = s_img.to(device), s_lbl.to(device)
            q_img, q_lbl = q_img.to(device), q_lbl.to(device)

            logits = model(s_img, s_lbl, q_img, n_way)
            loss = crit(logits, q_lbl)

            optim.zero_grad(); loss.backward(); optim.step()
            loss_sum += loss.item()
            acc_sum  += (logits.argmax(1) == q_lbl).float().mean().item()

        # ---- 验证阶段 ----
        model.eval()
        val_sampler = EpisodeSampler(val_data, val_lbls,
                                     n_way, n_support, n_query, n_episodes)
        val_acc = 0.0
        with torch.no_grad():
            for s_img, s_lbl, q_img, q_lbl in val_sampler:
                s_img, s_lbl = s_img.to(device), s_lbl.to(device)
                q_img, q_lbl = q_img.to(device), q_lbl.to(device)
                logits = model(s_img, s_lbl, q_img, n_way)
                val_acc += (logits.argmax(1) == q_lbl).float().mean().item()
        val_acc /= n_episodes
        print(f"  train_loss={loss_sum/n_episodes:.4f}  "
              f"train_acc={acc_sum/n_episodes:.4f}  val_acc={val_acc:.4f}")

        if val_acc > best:
            best = val_acc
            torch.save(model.state_dict(), 'best_protonet.pt')

# ---- 用随机数据跑一遍，确认流程通了 ----
if __name__ == '__main__':
    num_classes, samples_per_class, img_size = 64, 600, 84
    all_data = np.random.randn(num_classes * samples_per_class,
                               3, img_size, img_size).astype(np.float32)
    all_labels = np.repeat(np.arange(num_classes), samples_per_class)

    train_classes = int(num_classes * 0.8)
    train_mask = all_labels < train_classes
    val_mask   = all_labels >= train_classes

    encoder = ProtoNetEncoder()
    model = PrototypicalNetwork(encoder)
    train(model, all_data[train_mask], all_labels[train_mask],
          all_data[val_mask], all_labels[val_mask],
          n_way=5, n_support=5, n_query=15, num_epochs=10)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

import torch
import torch.nn as nn

class WeightedProtoNet(nn.Module):
    """基于支持集估计每类方差的原型网络。"""
    def __init__(self, encoder, eps=1e-3):
        super().__init__()
        self.encoder = encoder
        self.eps = eps  # 方差下限，避免 K=1 时爆炸

    def class_stats(self, support_emb, support_lbls, n_way):
        prototypes, inv_var = [], []
        for c in range(n_way):
            mask = (support_lbls == c)
            x_c  = support_emb[mask]                           # (K, D)
            mu_c = x_c.mean(dim=0)                             # (D,)
            # 各向同性每类方差：对特征维度取平均。
            if x_c.size(0) > 1:
                var_c = ((x_c - mu_c) ** 2).mean().clamp_min(self.eps)
            else:
                # K=1 情况：无方差信息，回退到全局先验。
                var_c = torch.tensor(1.0, device=mu_c.device)
            prototypes.append(mu_c)
            inv_var.append(1.0 / var_c)
        return torch.stack(prototypes), torch.stack(inv_var)   # (N,D), (N,)

    def forward(self, s_imgs, s_lbls, q_imgs, n_way):
        s_emb = self.encoder(s_imgs)
        q_emb = self.encoder(q_imgs)
        protos, inv_var = self.class_stats(s_emb, s_lbls, n_way)

        # 加权平方距离：-||q - mu_c||^2 / (2 sigma_c^2) - 0.5 d log sigma_c^2
        sq = ((q_emb.unsqueeze(1) - protos.unsqueeze(0)) ** 2).sum(dim=-1)  # (Q, N)
        log_norm = -0.5 * q_emb.size(1) * torch.log(1.0 / inv_var)          # (N,)
        return -0.5 * sq * inv_var + log_norm                               # logits

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

import torch
import torch.nn as nn
from torch.func import functional_call

class TinyMLP(nn.Module):
    def __init__(self, hidden=40):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(1, hidden), nn.ReLU(),
            nn.Linear(hidden, hidden), nn.ReLU(),
            nn.Linear(hidden, 1),
        )
    def forward(self, x):
        return self.net(x)

def sample_task(K=10):
    omega = torch.empty(1).uniform_(0.5, 2.0).item()
    phi   = torch.empty(1).uniform_(0.0, 6.2832).item()
    x = torch.empty(K, 1).uniform_(-5.0, 5.0)
    y = torch.sin(omega * x + phi)
    return x, y

def inner_step(model, params, x, y, alpha):
    """在支持集损失上执行一步 SGD，返回新参数（非原地更新）。"""
    pred = functional_call(model, params, (x,))
    loss = ((pred - y) ** 2).mean()
    grads = torch.autograd.grad(loss, params.values(), create_graph=True)
    return {name: p - alpha * g for (name, p), g in zip(params.items(), grads)}

def outer_step(model, meta_params, meta_optim, K=10, n_inner=5, alpha=0.01,
               n_tasks=4):
    """对 n_tasks 个任务的查询损失求和，执行一步元梯度更新。"""
    meta_optim.zero_grad()
    meta_loss = 0.0
    for _ in range(n_tasks):
        x_s, y_s = sample_task(K)
        x_q, y_q = sample_task(K)   # 新采样 -> 查询集
        params = meta_params
        for _ in range(n_inner):
            params = inner_step(model, params, x_s, y_s, alpha)
        pred_q = functional_call(model, params, (x_q,))
        meta_loss = meta_loss + ((pred_q - y_q) ** 2).mean()
    meta_loss.backward()         # 此处计算二阶梯度
    meta_optim.step()
    return meta_loss.item() / n_tasks

# --- 训练 ---
model = TinyMLP()
meta_params = {n: p.clone().detach().requires_grad_(True)
               for n, p in model.named_parameters()}
meta_optim = torch.optim.Adam(meta_params.values(), lr=1e-3)

for it in range(10000):
    loss = outer_step(model, meta_params, meta_optim)
    if it % 1000 == 0:
        print(f"iter {it}: meta-loss={loss:.4f}")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

import torch
import torch.nn as nn

class TaskConditionalMetric(nn.Module):
    """根据支持集摘要生成每 episode 的马氏矩阵。"""
    def __init__(self, embed_dim, hidden=128, rank=None):
        super().__init__()
        self.D = embed_dim
        self.rank = rank or embed_dim
        # 摘要：每类（均值，偏差平方均值）-> 2*D 个标量。
        # 超网络通过对类平均聚合，再映射到 L。
        self.hyper = nn.Sequential(
            nn.Linear(2 * embed_dim, hidden), nn.ReLU(),
            nn.Linear(hidden, embed_dim * self.rank),
        )

    def summary(self, support_emb, support_lbls, n_way):
        feats = []
        for c in range(n_way):
            x_c  = support_emb[support_lbls == c]
            mu   = x_c.mean(dim=0)
            var  = ((x_c - mu) ** 2).mean(dim=0) if x_c.size(0) > 1 else torch.zeros_like(mu)
            feats.append(torch.cat([mu, var], dim=-1))
        return torch.stack(feats).mean(dim=0)            # (2D,)

    def forward(self, support_emb, support_lbls, n_way):
        summary = self.summary(support_emb, support_lbls, n_way)
        L = self.hyper(summary).view(self.D, self.rank)
        M = L @ L.T + 1e-3 * torch.eye(self.D, device=L.device)   # 半正定，良态
        return M

def mahalanobis_logits(query_emb, prototypes, M):
    diff = query_emb.unsqueeze(1) - prototypes.unsqueeze(0)       # (Q, N, D)
    # (Q, N, D) @ (D, D) -> (Q, N, D)；再与 diff 点积 -> (Q, N)
    return -torch.einsum('qnd,de,qne->qn', diff, M, diff)