1
# MCDformer-main/exp_config/zhenhua/rml16a.py（load_rawdata 精简）
2
Set = pickle.load(open('RML2016.10a_dict.pkl', 'rb'), encoding='bytes')
3
snrs, mods = ...                    # 键里取出 SNR 列表与调制名列表
4
for mod in mods:
5
    for snr in snrs:
6
        Signals.append(Set[(mod, snr)])       # 每格 (N, 2, 128)

1
# meta-transformer-amc-main/data/dataset.py（节选）
2
self.data = h5py.File("GOLD_XYZ_OSC.0001_1024.hdf5", 'r')
3
self.iq      = self.data['X']     # (2555904, 1024, 2) 注意: (时间, I/Q) 布局!
4
self.onehot  = self.data['Y']     # (2555904, 24) 调制类型 one-hot
5
self.snr     = np.squeeze(self.data['Z'])   # (2555904,) 每样本的 SNR
6

7
x = self.iq[item].transpose()      # → (2, 1024)：转成 (I/Q, 时间)
8
x = np.expand_dims(x, axis=1)      # → (1, 2, 1024)：补出"图像通道"维

1
# meta-transformer-amc-main/models/vit.py（节选，注释为本教程添加）
2
class PatchEmbedding(nn.Module):
3
    def __init__(self, in_channels, embed_dim, patch_size):
4
        super().__init__()
5
        # in_channels=1: IQ 两行合并视为"1 通道图像"
6
        # patch_size=(2,16): 高 2 宽 16 —— 恰好覆盖 I/Q 两行 × 16 个时间点
7
        self.proj = nn.Conv2d(in_channels, embed_dim,
8
                              kernel_size=patch_size, stride=patch_size)
9

10
    def forward(self, x):
11
        x = self.proj(x)                     # (B,1,2,1024) → (B,36,1,64)
12
        x = x.flatten(2).transpose(1, 2)     # → (B,64,36)：64 个 token
13
        return x

1
Conv2d(1→36, k=(2,16), s=(2,16)):
2
  输出高 = (2-2)/2 + 1  = 1      ← I/Q 两行被"吃"进同一个 patch
3
  输出宽 = (1024-16)/16 + 1 = 64  ← 64 个时间窗
4
  → (B, 36, 1, 64)
5
flatten(2).transpose(1,2):
6
  → (B, 64, 36)                   ← 64 个 token，每个 36 维

1
# vit.py 节选
2
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
3
self.pos_embed = nn.Parameter(
4
    torch.zeros(1, 1 + (2 * 1024) // (2 * 16), embed_dim))   # 1+64=65 个位置
5

6
def forward(self, x):
7
    B, _, _, _ = x.shape
8
    x = self.patch_embed(x)                  # (B,64,36)
9
    x = torch.cat((self.cls_token.expand(B, -1, -1), x), dim=1)  # (B,65,36)
10
    x = x + self.pos_embed                   # 可学习位置编码
11
    x = self.pos_drop(x)
12
    x = self.blocks(x)                       # 8 层 ViTBlock
13
    x = self.norm(x[:, 0])                   # 取 cls token
14
    x = self.fc(x)                           # 分类头
15
    return x

1
class PatchEmbed(nn.Module):
2
    def __init__(self, patch_size=(2,8), in_chans=1, embed_dim=256, ...):
3
        self.num_patches = 128 // patch_size[1]     # 16 个 patch（L=128, patch 宽 8）
4
        self.proj = nn.Conv2d(in_chans, embed_dim,
5
                              kernel_size=patch_size, stride=patch_size, bias=True)
6
        self.norm = nn.BatchNorm2d(embed_dim)       # ★ 注意: 用 BatchNorm 而非 LayerNorm
7

8
    def forward(self, x):
9
        x = self.proj(x)                    # (B,1,2,128) → (B,256,1,16)
10
        x = torch.tanh(self.norm(x))        # ★ BN + tanh 压幅
11
        x = x.flatten(2).transpose(1, 2)    # → (B,16,256)
12
        return x

1
# MCDformer-main/models/MCDformer.py（节选，注释为本教程添加）
2
self.kernel_size_list = [(1,3), (2,3), (1,3), (1,3)]
3
self.conv_chan_list   = [1, 36, 64, 128, 256]     # 通道数逐层翻倍
4

5
class Conv_Block(nn.Module):
6
    def __init__(self, in_channel, out_channel, kernel_size):
7
        super().__init__()
8
        self.conv_block = nn.Sequential(
9
            nn.ZeroPad2d((1, 1, 0, 0)),            # 只补时间维(左右各 1)，保持长度
10
            nn.Conv2d(in_channel, out_channel, kernel_size=kernel_size),
11
            nn.ReLU(inplace=True),
12
            nn.BatchNorm2d(out_channel))
13

14
def forward(self, x):
15
    x = x.unsqueeze(1)         # (B,2,128) → (B,1,2,128)
16
    x = self.FDDM(x)           # 频域去噪（第 6 章）
17
    x = self.Conv_stem(x)      # 4 层 Conv_Block
18
    x = x.squeeze(2)           # (B,256,1,128) → (B,256,128)
19
    x = self.block(x)          # Transformer Block(dim=128) ← 通道当 token!

1
输入        (B, 1, 2, 128)    1 通道: I/Q 两行
2
Conv(1,3)  (B, 36, 2, 128)    pad(1,1) 保持宽 128；高不变
3
Conv(2,3)  (B, 64, 1, 128)    ★ 核高=2: I/Q 两行融合成 1 行!
4
Conv(1,3)  (B,128, 1, 128)
5
Conv(1,3)  (B,256, 1, 128)    通道 256，时间 128
6
squeeze(2) (B,256, 128)       → 送入 Block(dim=128)

1
import torch
2
import torch.nn as nn
3

4
class IQPatchEmbed(nn.Module):
5
    """可复用的 IQ 信号 patch 嵌入（Meta-Transformer 风格，补上位置编码裁剪）。
6
    用法:
7
        emb = IQPatchEmbed(patch_len=16, embed_dim=36)
8
        x = torch.randn(4, 2, 1024)          # (B, 2, L) 裸 IQ
9
        tokens = emb(x)                      # (B, 65, 36)  64 patch + cls
10
    """
11

12
    def __init__(self, patch_len: int = 16, embed_dim: int = 36,
13
                 use_cls: bool = True, dropout: float = 0.1):
14
        super().__init__()
15
        self.use_cls = use_cls
16
        self.patch_len = patch_len
17
        # 核高固定 2：一个 patch 必然覆盖 I/Q 两行（2.1 节设计动机 1）
18
        self.proj = nn.Conv2d(1, embed_dim,
19
                              kernel_size=(2, patch_len), stride=(2, patch_len))
20
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) \
21
            if use_cls else None
22
        # 位置编码惰性创建，长度按实际输入自适应（4.4 节思想）
23
        self.pos_embed = nn.Parameter(torch.zeros(1, 1, embed_dim))
24
        self.drop = nn.Dropout(dropout)
25

26
    def forward(self, x: torch.Tensor) -> torch.Tensor:
27
        # x: (B, 2, L) —— 裸 IQ（调用方不必预先 unsqueeze）
28
        B, _, L = x.shape
29
        x = self.proj(x.unsqueeze(1))                # (B, E, 1, N)
30
        x = x.flatten(2).transpose(1, 2)             # (B, N, E)
31
        if self.use_cls:
32
            x = torch.cat([self.cls_token.expand(B, -1, -1), x], dim=1)
33
        # 位置编码按 token 数裁剪/扩展（对比 CTDNN 的 pos_embed[:, :N]）
34
        if self.pos_embed.size(1) != x.size(1):
35
            if self.pos_embed.size(1) < x.size(1):
36
                # 不够长：线性插值扩展（ViT 处理不同分辨率的官方做法）
37
                pos = self.pos_embed.transpose(1, 2)                     # (1,E,T)
38
                pos = torch.nn.functional.interpolate(pos, size=x.size(1),
39
                                                      mode="linear")
40
                self.pos_embed = nn.Parameter(pos.transpose(1, 2))
41
            else:
42
                self.pos_embed = nn.Parameter(self.pos_embed[:, : x.size(1)])
43
        return self.drop(x + self.pos_embed)
44

45

46
# ---- 冒烟测试 ----
47
emb = IQPatchEmbed()
48
print(emb(torch.randn(4, 2, 1024)).shape)     # (4, 65, 36)  1024 长
49
print(emb(torch.randn(4, 2, 256)).shape)      # (4, 17, 36)  256 长自动适配!

1
N = 256   ← "token" = 256 个卷积通道（每种通道 = 一种时域滤波器输出）
2
C = 128   ← 每个 token 的"特征" = 128 长的时域响应

1
def forward(self, x):
2
    B, N, C = x.shape                       # (B, 256, 128)：N=通道数, C=时间长度
3
    q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
4
    # ... q/k/v 各 (B, H, N, d_h) = (B, 2, 256, 64)
5
    attn = (q @ k.transpose(-2, -1)) * self.scale   # (B, 2, 256, 256) ★ 通道×通道!
6
    attn = attn.softmax(dim=-1)
7
    x = (attn @ v).transpose(1, 2).reshape(B, N, C)

1
stem 输出      (B, 256, 1, 128)
2
squeeze(2)     (B, 256, 128)       256 token × 128 维
3
Block(dim=128) QKV 投影: 128→128   每个"通道 token"的时间响应被投影
4
  attn         (B, H, 256, 256)    ★ 通道间交互
5
  MLP          (B, 256, 128)       沿 128 时间维逐 token 混合
6
ConvLayer(256) Conv1d + MaxPool(3,2,1): (B, 256, 64)   时间维减半
7
Block2(dim=64) 同构, 特征维 64
8
x[:, :, -1:]   (B, 256, 1)         取时间维末位
9
classifier     (B, 256) → (B, 11)

1
              ┌──────────────┬──────────────────┐
2
              │ 标准用法      │ MCDformer 用法    │
3
──────────────┼──────────────┼──────────────────┤
4
注意力（跨轴） │ 时间步 ↔ 时间步 │ 通道 ↔ 通道        │
5
MLP（轴内）   │ 通道内        │ 时间轴内          │

1
class FeaFusionModule(nn.Module):
2
    def __init__(self, num_attention_heads, input_size, hidden_size):
3
        ...
4
        self.key_layer   = nn.Linear(input_size, hidden_size)
5
        self.query_layer = nn.Linear(input_size, hidden_size)
6
        self.value_layer = nn.Linear(input_size, hidden_size)
7
        self.dropout = nn.Dropout(0.5)
8

9
    def trans_to_multiple_heads(self, x):
10
        new_size = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
11
        x = x.view(new_size)
12
        return x.permute(0, 2, 1, 3)              # (B, H, N, d_h)
13

14
    def forward(self, x):
15
        key = self.key_layer(x); query = self.query_layer(x); value = self.value_layer(x)
16
        key_heads = self.trans_to_multiple_heads(key)      # (B,H,256,64)
17
        # ...
18
        attention_scores = torch.matmul(query_heads, key_heads.permute(0, 1, 3, 2))
19
        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
20
        attention_probs = F.softmax(attention_scores, dim=-1)
21
        context = torch.matmul(attention_probs, value_heads)      # (B,H,256,64)
22
        context = context.contiguous().view(shape[0], -1, shape[-1])  # (B,512,64)!
23
        return context

1
class ChannelAttnBlock(nn.Module):
2
    """MCDformer 风格转置块：输入 (B, C, L)，token = 通道 C，特征维 = 时间 L。
3
    与教学文档 3.6 节 EncoderBlock 的唯一区别：不转置、直接吃 (B,C,L)。"""
4

5
    def __init__(self, sig_len: int, n_heads: int = 2, mlp_ratio: float = 1.0,
6
                 dropout: float = 0.3):
7
        super().__init__()
8
        assert sig_len % n_heads == 0
9
        self.norm1 = nn.LayerNorm(sig_len)              # 沿时间维归一化!
10
        self.attn = nn.MultiheadAttention(sig_len, n_heads,   # embed_dim = 时间长度
11
                                          dropout=dropout, batch_first=True)
12
        self.norm2 = nn.LayerNorm(sig_len)
13
        hidden = int(sig_len * mlp_ratio)
14
        self.ffn = nn.Sequential(nn.Linear(sig_len, hidden), nn.GELU(),
15
                                 nn.Dropout(dropout), nn.Linear(hidden, sig_len))
16
        self.drop = nn.Dropout(dropout)
17

18
    def forward(self, x: torch.Tensor) -> torch.Tensor:
19
        # x: (B, C, L) —— 卷积 stem 的输出天然是这个布局，直接吃
20
        x = x + self.drop(self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0])
21
        x = x + self.drop(self.ffn(self.norm2(x)))
22
        return x
23

24

25
# ---- 冒烟测试 ----
26
blk = ChannelAttnBlock(sig_len=128)
27
x = torch.randn(4, 256, 128)               # 256 通道 × 128 时间
28
y = blk(x)
29
print(y.shape)                             # (4, 256, 128)
30
# 注意力在 256 个通道间交互；LayerNorm 沿 128 维（时间）归一化

1
# meta-transformer-amc-main/data/dataset.py（FewShotDataset 节选）
2
if self.padding == 'self_duplicate':
3
    num_dup = 1024 // L
4
    query_set = [np.concatenate([self.iq[i].transpose()[:, :L]
5
                                 for _ in range(num_dup)], axis=1) ...]  # 重复拼接 4 次
6
elif self.padding == 'zero':
7
    query_set = [np.concatenate((sig, np.zeros((2, 1024-L))), axis=1) ...]  # 补零

1
# CTDNN.py 节选
2
x = x + self.pos_embed[:, : x.size(1), :]   # 按实际 token 数裁剪

1
pos = self.pos_embed.transpose(1, 2)                    # (1, E, T)
2
pos = F.interpolate(pos, size=new_len, mode="linear")   # 插值到新长度

1
def adapt_length(x, target_len, mode="duplicate"):
2
    """把 (B, 2, L) 的短帧适配到 target_len。三种策略：
3
    'duplicate': 重复拼接（频谱内插，无损）—— Meta-Transformer 默认
4
    'zero'     : 补零（谱峰展宽）         —— 不推荐但实现简单
5
    'crop'     : 只取前 target_len（L > target_len 时用）
6
    """
7
    B, _, L = x.shape
8
    if L == target_len:
9
        return x
10
    if L < target_len and mode == "duplicate":
11
        k = target_len // L
12
        rest = x[:, :, : target_len - k * L]
13
        return torch.cat([x.repeat(1, 1, k), rest], dim=2)
14
    if L < target_len and mode == "zero":
15
        return F.pad(x, (0, target_len - L))
16
    return x[:, :, :target_len]
17

18

19
x = torch.randn(4, 2, 256)
20
print(adapt_length(x, 1024, "duplicate").shape)   # (4, 2, 1024)
21
print(adapt_length(x, 1024, "zero").shape)        # (4, 2, 1024)
22
print(adapt_length(x, 128, "crop").shape)         # (4, 2, 128)

1
# A. Pre-LN（vit.py、MCDformer.py 同款；教学文档 3.6 节）—— 主流推荐
2
x = x + attn(norm1(x))
3
x = x + ffn(norm2(x))
4

5
# B. Post-LN（原论文）—— 现代实现已少用
6
x = norm1(x + attn(x))
7
x = norm2(x + ffn(x))
8

9
# C. CTDNN 变体（CTDNN.py 实际代码）—— 混合形态
10
x = norm1(x + attn(x))       # 残差后归一，但 attn 的输入没有归一化
11
x = norm2(x + dropout(ffn(x)))

1
# vit.py：ViTBlock(embed_dim, num_heads, mlp_dim) → MLP(embed_dim, int(embed_dim * mlp_ratio), ...)
2
# 即隐层 = 36 × 32 = 1152（而不是字面意义的 32）！

1
class ConfigurableBlock(nn.Module):
2
    """把本章全部自由度做成参数：归一化排布、mlp 倍率、dropout 强度。"""
3

4
    def __init__(self, d, n_heads, mlp_ratio=4.0, dropout=0.1,
5
                 norm_style="pre", act=nn.GELU):
6
        super().__init__()
7
        assert norm_style in ("pre", "post", "ctdnn")
8
        self.style = norm_style
9
        self.n1 = nn.LayerNorm(d)
10
        self.attn = nn.MultiheadAttention(d, n_heads, dropout=dropout, batch_first=True)
11
        self.n2 = nn.LayerNorm(d)
12
        self.ffn = nn.Sequential(nn.Linear(d, int(d * mlp_ratio)), act(),
13
                                 nn.Dropout(dropout), nn.Linear(int(d * mlp_ratio), d))
14
        self.drop = nn.Dropout(dropout)
15

16
    def forward(self, x):
17
        a = lambda t: self.drop(self.attn(t, t, t)[0])   # 自注意力支路
18
        f = lambda t: self.drop(self.ffn(t))
19
        if self.style == "pre":
20
            x = x + a(self.n1(x)); x = x + f(self.n2(x))
21
        elif self.style == "post":
22
            x = self.n1(x + a(x)); x = self.n2(x + f(x))
23
        else:  # ctdnn
24
            x = self.n1(x + a(x)); x = self.n2(x + f(x))
25
            x = self.n2(x)      # CTDNN 变体里 MLP 输入未归一——此简化仅示意
26
        return x

1
import numpy as np
2
import matplotlib.pyplot as plt   # 仅本节演示用
3

4
t = np.arange(128) / 128
5
s = np.sin(2 * np.pi * 5 * t) + 0.5 * np.sin(2 * np.pi * 20 * t)   # 5Hz + 20Hz 混合
6
S = np.fft.fft(s)
7
freqs = np.fft.fftfreq(128, 1 / 128)
8
print("频谱峰值位置:", freqs[np.argsort(-np.abs(S))[:4]])   # [5, -5, 20, -20] ✓

1
# MCDformer-main/models/MCDformer.py（按实际张量形状整理，注释为本教程添加）
2
class TinyMLP(nn.Module):
3
    """学习到的频域滤波器：L → L/4 → L，Tanh 把增益限幅在 (-1, 1)。"""
4
    def __init__(self, N):
5
        super().__init__()
6
        self.mlp = nn.Sequential(
7
            nn.Linear(N, N // 4), nn.ReLU(inplace=True),
8
            nn.Linear(N // 4, N), nn.Tanh())       # ★ Tanh: 增益有界
9

10
class FrequencyDomainDenoisingModule(nn.Module):
11
    def __init__(self, N):          # N = 信号长度 128
12
        super().__init__()
13
        self.mlp = TinyMLP(N)
14

15
    def forward(self, x):
16
        # x: (B, 1, 2, L) —— 行 0 = I，行 1 = Q
17
        x_init = x.clone()                          # 留底做残差
18
        r = x[:, :, 0, :] + 1j * x[:, :, 1, :]      # ① I + jQ = 复信号 (B,1,L)
19
        R = torch.fft.fft(r, dim=-1)                # ② 频谱 (B,1,L)
20
        R_stack = torch.stack([R.real, R.imag], dim=2)   # ③ (B,1,2,L) 实虚分两行
21
        h = self.mlp(R_stack)                       # ④ 学习频域增益 (B,1,2,L)
22
        R_hat = h[:, :, 0, :] * R_stack[:, :, 0, :] \
23
              + 1j * (h[:, :, 1, :] * R_stack[:, :, 1, :])   # ⑤ 实/虚谱分别加权
24
        x_hat = torch.fft.ifft(R_hat, dim=-1).real.unsqueeze(2)   # ⑥ 回时域 (B,1,1,L)
25
        return x_hat + x_init                       # ⑦ 残差连接

1
# AMC_Net.py 节选
2
class AdaCorrModule(nn.Module):
3
    def __init__(self, N):
4
        self.Im = TinyMLP(N)     # 虚部谱专用滤波器
5
        self.Re = TinyMLP(N)     # 实部谱专用滤波器
6

7
    def forward(self, x):
8
        x_init = copy.deepcopy(x)
9
        x = torch.fft.fft(x, dim=-1)          # 对两行分别 FFT
10
        X_re, X_im = torch.real(x), torch.imag(x)
11
        h_re = self.Re(X_re); h_im = self.Im(X_im)      # 实虚分治!
12
        x = torch.mul(h_re, X_re) + 1j * torch.mul(h_im, X_im)
13
        x = torch.real(torch.fft.ifft(x, dim=-1))       # 只取实部
14
        x = x + x_init
15
        return x

1
class FDDM(nn.Module):
2
    """独立可复用的频域去噪模块。use_complex=False 时退化为实数 FFT 版本
3
    （对实值信号任务可直接用）。"""
4

5
    def __init__(self, sig_len: int, hidden_ratio: int = 4):
6
        super().__init__()
7
        self.filter = nn.Sequential(
8
            nn.Linear(sig_len, sig_len // hidden_ratio), nn.ReLU(),
9
            nn.Linear(sig_len // hidden_ratio, sig_len), nn.Tanh())
10

11
    def forward(self, x: torch.Tensor) -> torch.Tensor:
12
        # x: (B, 1, 2, L) —— I/Q 两行
13
        r = x[:, :, 0, :] + 1j * x[:, :, 1, :]
14
        R = torch.fft.fft(r, dim=-1)
15
        R_stack = torch.stack([R.real, R.imag], dim=2)
16
        h = self.filter(R_stack)
17
        R_hat = h[:, :, 0, :] * R_stack[:, :, 0, :] \
18
              + 1j * (h[:, :, 1, :] * R_stack[:, :, 1, :])
19
        return torch.fft.ifft(R_hat, dim=-1).real.unsqueeze(2) + x
20

21

22
fddm = FDDM(128)
23
x = torch.randn(4, 1, 2, 128)
24
print(fddm(x).shape)                # (4, 1, 2, 128) —— 形状不变
25
# 残差保证: 最坏情况模块学成恒等（把 filter 权重冻结为 0 即可验证）

1
输入 → [物理先验模块(FFT/小波/谱图/均衡…)，必须带残差] → [卷积 stem] → [Transformer] → 任务头

1
# meta-transformer-amc-main/models/proto.py（proto_train 核心，注释为本教程添加）
2
def proto_train(self, sample):
3
    n_way = len(sample.keys())            # 本 episode 有 5 个类
4
    n_support = self.config['num_support']   # 每类 5 个有标签样本
5
    n_query = self.config['num_query']       # 每类 10 个查询样本
6

7
    # ---- ① 把 episode 字典摊平成两个张量 ----
8
    # support: (n_way×n_support, 2, L)；query: (n_way×n_query, 2, L)
9
    x_support = np.vstack([np.array([np.array(iq) for iq in sample[label]['support']])
10
                           for label in sample.keys()])
11
    x_query   = np.vstack([... 同理 ...])
12
    x_support = torch.from_numpy(x_support).cuda(0)
13
    x_query   = torch.from_numpy(x_query).cuda(0)
14

15
    # ---- ② 编码：同一个 encoder 前向两次 ----
16
    z_support = self.encoder.forward(x_support)      # (75, d)
17
    z_query   = self.encoder.forward(x_query)        # (50, d)
18

19
    # ---- ③ 原型 = support 嵌入按类求均值 ----
20
    z_proto = z_support.view(n_way, n_support, -1).mean(1)   # (5, d)
21

22
    # ---- ④ 距离 → 概率 ----
23
    dists = torch.cdist(z_query, z_proto)            # (50, 5) 欧氏距离
24
    log_p_y = F.log_softmax(-dists, dim=1)           # 负距离 softmax = 概率
25

26
    # ---- ⑤ 交叉熵 + 准确率 ----
27
    target_inds = torch.arange(0, n_way).view(n_way, 1, 1) \
28
                      .expand(n_way, n_query, 1).long()     # 每类的真值下标
29
    loss = -log_p_y.gather(2, target_inds).squeeze().view(-1).mean()
30
    _, y_hat = log_p_y.max(2)
31
    acc = torch.eq(y_hat, target_inds.squeeze()).float().mean()
32
    return loss, {'loss': ..., 'acc': ...}

1
# proto.py load_protonet_vit
2
encoder = ViT(..., num_classes=model_params["num_classes"])  # fc 还在!
3
# proto_train 里: z_support = self.encoder.forward(x_support)  # 含 fc 输出

1
# dataset.py 监督路径做了 expand_dims:
2
x = np.expand_dims(x, axis=1)      # (2,1024) → (1,2,1024) ✓
3
# 但 FewShotDataset 的 support/query 没有这一步 ✗

1
assert x.dim() == 4 and x.size(1) == 1, f"期望 (B,1,2,L)，实际 {x.shape}"

1
def proto_loss(encoder, x_support, y_support, x_query, y_query):
2
    """原型网络损失（7.2 节公式的完整实现）。
3
    encoder: 任意"输入 → (B,d) 嵌入"的模型（Transformer/CNN 均可）
4
    x_support: (n_way×n_sup, ...)  y_support: 类标签（可非 0 起始）
5
    返回 (loss, acc)。"""
6
    z_s = encoder(x_support)
7
    z_q = encoder(x_query)
8
    classes = torch.unique(y_support)                    # 本 episode 实际出现的类
9
    protos = torch.stack([z_s[y_support == c].mean(0) for c in classes])  # (n_way, d)
10
    remap = {c.item(): i for i, c in enumerate(classes)}  # 类 id → 0..n_way-1
11
    y_q = torch.tensor([remap[v] for v in y_query.tolist()])
12
    dists = torch.cdist(z_q, protos)
13
    loss = F.cross_entropy(-dists, y_q)
14
    acc = (-dists).argmax(1).eq(y_q).float().mean()
15
    return loss, acc

1
监督学习:   [数据 → 模型 → 24 类 logits]   每 batch 学"这 24 类怎么分"
2
元学习:     [随机抽 5 类 → 每类 5 样本算原型 → 10 查询样本分类]
3
            每个 episode 都是一个新的 5 类小任务 → 学"如何从小样本分类"

1
# data/dataset.py FewShotDataset.__init__（节选）
2
self.label_indices = {label: [i for i, x in enumerate(self.label_list) if x == label]
3
                      for label in self.labels}     # 每个类 → 它所有样本的下标
4
self.num_support = self.config["num_support"]       # 5
5
self.num_query   = self.config["num_query"]         # 10
6
self.num_episode = len(self.snr) // ((self.num_support + self.num_query) * len(self.labels))
7

8
# __getitem__（节选）: idx = 第几个 episode
9
for label in self.labels:                            # 遍历本数据集的所有类
10
    label_indices = self.label_indices[label]
11
    support_indices = random.sample(label_indices, self.num_support)   # 抽 5 个 support
12
    sample[label]['support'] = [self.iq[i].transpose()[:, :L] for i in support_indices]
13
    query_indices = list(set(label_indices) - set(support_indices))    # 剩余里抽 query
14
    query_indices = random.sample(query_indices, self.num_query)       # 10 个
15
    sample[label]['query'] = [...]

1
# runner/train.py meta_train（节选）
2
train_dataloader = DATA.DataLoader(train_data, batch_size=1, shuffle=True)  # ★ batch=1!
3
for epoch in range(self.model_params["epoch"]):
4
    for episode, sample in enumerate(tqdm.tqdm(train_dataloader)):
5
        self.optimizer.zero_grad()
6
        loss, output = self.net.proto_train(sample)   # 一个 episode = 一个"任务"
7
        loss.backward()
8
        self.optimizer.step()

1
① HDF5 原始样本          X[i]: (1024, 2)            数据集布局: 时间 × I/Q
2
② dataset transpose      → (2, 1024)                转成 I/Q × 时间
3
③ FewShotDataset 抽 episode
4
      sample = {类0: {support:[5 个 (2,1024)], query:[10 个]}, ... 共 5 类}
5
④ proto_train 摊平
6
      x_support: (75, 2, 1024)    x_query: (50, 2, 1024)
7
⑤ [设计意图] 补通道维     (75, 1, 2, 1024)           ViT 入口期望 4D（7.4 坑 2）
8
⑥ PatchEmbedding(2×16)  (75, 36, 1, 64) → (75, 64, 36)   64 个 IQ patch token
9
⑦ + cls_token           (75, 65, 36)               cls 占位 0
10
⑧ + pos_embed(可学习)    (75, 65, 36)               长度 65 固定
11
⑨ 8 × ViTBlock          (75, 65, 36)               Pre-LN 自注意力 ×8 层
12
⑩ norm(x[:,0])          (75, 36)                   cls 嵌入
13
⑪ fc（repo 实际含此层）  (75, 24)                   ★ 7.4 坑 1: 嵌入实为 24 维 logits
14
⑫ 原型                   (5, 24)                    support 按类取均值
15
⑬ cdist                  (50, 5)                    查询 × 原型距离矩阵
16
⑭ log_softmax(-dists)    (50, 5)                    → 概率
17
⑮ 交叉熵 + 准确率         标量                       与真值 (50,) 比较

1
① 输入 IQ               (B, 2, 128)
2
② unsqueeze(1)          (B, 1, 2, 128)              视作 1 通道"矮图像"
3
③ FDDM 频域去噪         (B, 1, 2, 128)              I+jQ → FFT → MLP 增益 → IFFT → +残差
4
④ Conv_Block(1→36)      (B, 36, 2, 128)             核(1,3) 只沿时间
5
⑤ Conv_Block(36→64)     (B, 64, 1, 128)             ★ 核(2,3): I/Q 融合成 1 行
6
⑥ Conv_Block(64→128)    (B, 128, 1, 128)
7
⑦ Conv_Block(128→256)   (B, 256, 1, 128)            通道 256
8
⑧ squeeze(2)            (B, 256, 128)               256 token × 128 维 ← 转置语义!
9
⑨ Block(dim=128)        (B, 256, 128)               注意力矩阵 (B,H,256,256): 通道交互
10
⑩ ConvLayer(256)        (B, 256, 64)                Conv1d(循环填充)+ELU+MaxPool: 时间减半
11
⑪ Block2(dim=64)        (B, 256, 64)                第二级通道注意力
12
⑫ x[:, :, -1:]          (B, 256, 1)                 末位读出
13
⑬ classifier            (B, 11)                     Linear(256→256)→PReLU→Linear(256→11)

1
import math
2
import torch
3
import torch.nn as nn
4
import torch.nn.functional as F
5

6
torch.manual_seed(0)
7

8
# ---- 设备选择：默认 GPU ----
9
# set_default_device 让之后创建的所有张量/模块自动落在 GPU 上，正文代码无需再写 .to()/.cuda()
10
torch.set_default_device('cuda')
11
print("默认设备: cuda ——", torch.cuda.get_device_name(0))
12

13
# ================= 迷你 ViT 编码器（Meta-Transformer 风格） =================
14
class IQPatchEmbed(nn.Module):
15
    """2×16 patch：一个 token 同时含 I/Q 两行（第 2.1 节）。"""
16
    def __init__(self, patch_len=16, embed_dim=36):
17
        super().__init__()
18
        self.proj = nn.Conv2d(1, embed_dim,
19
                              kernel_size=(2, patch_len), stride=(2, patch_len))
20

21
    def forward(self, x):                     # (B,1,2,L) → (B,N,E)
22
        return self.proj(x).flatten(2).transpose(1, 2)
23

24

25
class MHA(nn.Module):
26
    """多头注意力（教学文档 3.4 节精简版）。"""
27
    def __init__(self, d, H):
28
        super().__init__()
29
        self.H, self.dh = H, d // H
30
        self.qkv = nn.Linear(d, 3 * d)
31
        self.proj = nn.Linear(d, d)
32

33
    def forward(self, x):
34
        B, N, d = x.shape
35
        qkv = self.qkv(x).reshape(B, N, 3, self.H, self.dh).permute(2, 0, 3, 1, 4)
36
        q, k, v = qkv[0], qkv[1], qkv[2]
37
        a = (q @ k.transpose(-2, -1)) * (self.dh ** -0.5)
38
        a = a.softmax(-1)
39
        return self.proj((a @ v).transpose(1, 2).reshape(B, N, d))
40

41

42
class Block(nn.Module):
43
    """Pre-LN 块（第 5.1 节 A 式）。"""
44
    def __init__(self, d, H):
45
        super().__init__()
46
        self.n1 = nn.LayerNorm(d); self.attn = MHA(d, H)
47
        self.n2 = nn.LayerNorm(d)
48
        self.ffn = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
49

50
    def forward(self, x):
51
        x = x + self.attn(self.n1(x))
52
        x = x + self.ffn(self.n2(x))
53
        return x
54

55

56
class MiniViTEncoder(nn.Module):
57
    """(B,2,L) → (B,E) 嵌入。★ 没有分类头——度量学习专用（第 7.2 节）。"""
58

59
    def __init__(self, L=1024, patch_len=16, embed_dim=36, n_heads=4, n_layers=2):
60
        super().__init__()
61
        n_patches = L // patch_len
62
        self.patch = IQPatchEmbed(patch_len, embed_dim)
63
        self.cls = nn.Parameter(torch.zeros(1, 1, embed_dim))         # cls token
64
        self.pos = nn.Parameter(torch.zeros(1, n_patches + 1, embed_dim))
65
        self.blocks = nn.ModuleList([Block(embed_dim, n_heads)
66
                                     for _ in range(n_layers)])
67
        self.norm = nn.LayerNorm(embed_dim)
68

69
    def encode(self, x):
70
        x = self.patch(x.unsqueeze(1))                     # (B,N,E)
71
        x = torch.cat([self.cls.expand(x.size(0), -1, -1), x], dim=1)
72
        x = x + self.pos
73
        for blk in self.blocks:
74
            x = blk(x)
75
        return self.norm(x[:, 0])                          # cls 嵌入 (B,E)
76

77

78
# ================= 数据: 5 类复指数 IQ 信号（模拟不同"调制"） =================
79
def make_iq(n_per_class=200, L=1024, spec=None, noise=0.3):
80
    """类别由 (频率周期数, 振幅) 区分，加复高斯噪声。"""
81
    if spec is None:
82
        spec = [(4, 1.0), (8, 1.0), (16, 1.0),   # 训练类
83
                (12, 1.0), (24, 1.0)]            # ★ 测试类: 元训练时从未见过
84
    t = torch.arange(L).float() / L * 2 * math.pi
85
    xs, ys = [], []
86
    for c, (f, amp) in enumerate(spec):
87
        phase = torch.rand(n_per_class, 1) * 2 * math.pi
88
        s = amp * torch.exp(1j * (f * t + phase))                     # (n, L)
89
        s = s + noise * (torch.randn(n_per_class, L)
90
                         + 1j * torch.randn(n_per_class, L))          # 噪声
91
        iq = torch.stack([s.real, s.imag], dim=1)                     # (n, 2, L)
92
        xs.append(iq)
93
        ys.append(torch.full((n_per_class,), c, dtype=torch.long))
94
    return torch.cat(xs), torch.cat(ys)
95

96

97
# ================= 元学习组件（第 7.5 / 8.2 节） =================
98
def make_episode(xs, ys, class_ids, n_sup=5, n_q=10):
99
    """从指定类别抽一个 episode: support 与 query 互斥。"""
100
    Sx, Sy, Qx, Qy = [], [], [], []
101
    for c in class_ids:
102
        idx = (ys == c).nonzero().flatten()
103
        idx = idx[torch.randperm(len(idx))]
104
        Sx.append(xs[idx[:n_sup]]);   Sy += [c] * n_sup
105
        Qx.append(xs[idx[n_sup:n_sup + n_q]]); Qy += [c] * n_q
106
    return (torch.cat(Sx), torch.tensor(Sy),
107
            torch.cat(Qx), torch.tensor(Qy))
108

109

110
def proto_loss(encoder, x_s, y_s, x_q, y_q):
111
    """原型网络损失（第 7.2 节公式）：原型 = support 嵌入均值，按最近原型分类。"""
112
    z_s, z_q = encoder.encode(x_s), encoder.encode(x_q)
113
    classes = torch.unique(y_s)
114
    protos = torch.stack([z_s[y_s == c].mean(0) for c in classes])   # (n_way, E)
115
    remap = {c.item(): i for i, c in enumerate(classes)}             # 类 id 重排
116
    y_q_r = torch.tensor([remap[v] for v in y_q.tolist()])
117
    dists = torch.cdist(z_q, protos)
118
    loss = F.cross_entropy(-dists, y_q_r)
119
    acc = (-dists).argmax(1).eq(y_q_r).float().mean()
120
    return loss, acc
121

122

123
# ================= 主流程 =================
124
def run_meta_demo():
125
    xs, ys = make_iq()
126
    train_classes = [0, 1, 2]        # 元训练只见过这 3 类
127
    test_classes = [3, 4]            # ★ 这 2 类整个训练过程从未出现
128

129
    enc = MiniViTEncoder(L=1024)
130
    opt = torch.optim.Adam(enc.parameters(), lr=1e-3)
131

132
    # ---- 元训练: 每个 episode 都是一个新的 3-way 小任务 ----
133
    for ep in range(200):
134
        loss_acc = 0.0
135
        for _ in range(4):                              # 每步 4 个 episode
136
            cids = torch.randperm(len(train_classes))[:3].tolist()
137
            xs_s, ys_s, xs_q, ys_q = make_episode(xs, ys, cids)
138
            loss, acc = proto_loss(enc, xs_s, ys_s, xs_q, ys_q)
139
            loss.backward()
140
            loss_acc += acc.item() / 4
141
        opt.step(); opt.zero_grad()
142
        if ep % 40 == 0 or ep == 199:
143
            print(f"episode {ep:3d}  训练 acc = {loss_acc:.3f}")
144

145
    # ---- 元测试: 对从未见过的 2 类做 2-way 5-shot（随机猜测 = 0.500）----
146
    accs = []
147
    enc.eval()
148
    for _ in range(100):
149
        xs_s, ys_s, xs_q, ys_q = make_episode(xs, ys, test_classes)
150
        with torch.no_grad():
151
            _, acc = proto_loss(enc, xs_s, ys_s, xs_q, ys_q)
152
        accs.append(acc.item())
153
    print(f"\n★ 未见类 2-way 5-shot 准确率: {sum(accs)/len(accs):.3f}  "
154
          f"(随机 = 0.500，监督学习 = 0.000)")
155

156

157
if __name__ == "__main__":
158
    run_meta_demo()

1
import math
2
import torch
3
import torch.nn as nn
4
import torch.nn.functional as F
5

6
torch.manual_seed(0)
7

8
# ---- 设备选择：默认 GPU ----
9
# set_default_device 让之后创建的所有张量/模块自动落在 GPU 上，正文代码无需再写 .to()/.cuda()
10
torch.set_default_device('cuda')
11
print("默认设备: cuda ——", torch.cuda.get_device_name(0))
12

13
# ================= MiniMCDformer: FDDM + Conv stem + 通道注意力 =================
14
class FDDM(nn.Module):
15
    """频域去噪（第 6.2 节）：FFT → 学习频域增益(Tanh 限幅) → IFFT → 残差。"""
16

17
    def __init__(self, sig_len: int):
18
        super().__init__()
19
        self.filter = nn.Sequential(
20
            nn.Linear(sig_len, sig_len // 4), nn.ReLU(),
21
            nn.Linear(sig_len // 4, sig_len), nn.Tanh())
22

23
    def forward(self, x):                      # (B,1,2,L)
24
        r = x[:, :, 0, :] + 1j * x[:, :, 1, :]             # 复信号
25
        R = torch.fft.fft(r, dim=-1)
26
        R_stack = torch.stack([R.real, R.imag], dim=2)     # (B,1,2,L)
27
        h = self.filter(R_stack)                           # 频域增益
28
        R_hat = h[:, :, 0, :] * R_stack[:, :, 0, :] \
29
              + 1j * (h[:, :, 1, :] * R_stack[:, :, 1, :])
30
        x_hat = torch.fft.ifft(R_hat, dim=-1).real.unsqueeze(2)
31
        return x_hat + x                                   # 残差安全网
32

33

34
class ConvBlock(nn.Module):
35
    """stem 卷积块（第 2.3 节）：只补时间维，保持长度。"""
36

37
    def __init__(self, cin, cout, k):
38
        super().__init__()
39
        self.net = nn.Sequential(
40
            nn.ZeroPad2d((1, 1, 0, 0)),
41
            nn.Conv2d(cin, cout, kernel_size=k), nn.ReLU(), nn.BatchNorm2d(cout))
42

43
    def forward(self, x):
44
        return self.net(x)
45

46

47
class ChannelAttnBlock(nn.Module):
48
    """通道注意力块（第 3.5 节）：输入 (B,C,L)，token=通道，特征维=时间。"""
49

50
    def __init__(self, sig_len: int, n_heads: int = 2, dropout: float = 0.2):
51
        super().__init__()
52
        self.n1 = nn.LayerNorm(sig_len)
53
        self.attn = nn.MultiheadAttention(sig_len, n_heads,      # embed_dim = 时间长度!
54
                                          dropout=dropout, batch_first=True)
55
        self.n2 = nn.LayerNorm(sig_len)
56
        self.ffn = nn.Sequential(nn.Linear(sig_len, sig_len), nn.GELU(),
57
                                 nn.Dropout(dropout), nn.Linear(sig_len, sig_len))
58

59
    def forward(self, x):
60
        x = x + self.attn(self.n1(x), self.n1(x), self.n1(x))[0]  # 通道间交互
61
        x = x + self.ffn(self.n2(x))                              # 时间维混合
62
        return x
63

64

65
class MiniMCDformer(nn.Module):
66
    """FDDM → Conv stem → 通道注意力 → 时间池化 → 通道注意力 → 末位读出。"""
67

68
    def __init__(self, L=128, num_classes=4, use_fddm=True):
69
        super().__init__()
70
        self.use_fddm = use_fddm
71
        self.fddm = FDDM(L)
72
        self.stem = nn.Sequential(
73
            ConvBlock(1, 32, (1, 3)),        # (B,32,2,L)
74
            ConvBlock(32, 64, (2, 3)),       # ★ 核高 2: I/Q 融合 → (B,64,1,L)
75
            ConvBlock(64, 128, (1, 3)),      # (B,128,1,L)
76
        )
77
        self.blk1 = ChannelAttnBlock(L)                        # 128 通道交互
78
        self.down = nn.Sequential(                             # 时间维减半
79
            nn.Conv1d(128, 128, 3, padding=1, padding_mode="circular"),
80
            nn.BatchNorm1d(128), nn.ELU(), nn.MaxPool1d(3, 2, 1))
81
        self.blk2 = ChannelAttnBlock(L // 2)
82
        self.head = nn.Sequential(nn.Linear(128, 64), nn.PReLU(),
83
                                  nn.Linear(64, num_classes))
84

85
    def forward(self, x):                      # (B,2,L)
86
        x = x.unsqueeze(1)                     # (B,1,2,L)
87
        if self.use_fddm:
88
            x = self.fddm(x)                   # 频域去噪（可关 = 消融）
89
        x = self.stem(x)                       # (B,128,1,L)
90
        x = x.squeeze(2)                       # (B,128,L) ← 通道即 token
91
        x = self.blk1(x)
92
        x = self.down(x)                       # (B,128,L/2)
93
        x = self.blk2(x)
94
        x = x[:, :, -1]                        # 末位读出 (B,128)
95
        return self.head(x)                    # (B,num_classes)
96

97

98
def make_iq_short(n_per_class=400, L=128, freqs=(4, 8, 16, 32), noise=0.8):
99
    """4 类不同频率复指数 + 强噪声（低 SNR 场景，考验去噪能力）。"""
100
    t = torch.arange(L).float() / L * 2 * math.pi
101
    xs, ys = [], []
102
    for c, f in enumerate(freqs):
103
        phase = torch.rand(n_per_class, 1) * 2 * math.pi
104
        s = torch.exp(1j * (f * t + phase))
105
        s = s + noise * (torch.randn(n_per_class, L) + 1j * torch.randn(n_per_class, L))
106
        iq = torch.stack([s.real, s.imag], dim=1)                 # (n,2,L)
107
        xs.append(iq); ys.append(torch.full((n_per_class,), c, dtype=torch.long))
108
    return torch.cat(xs), torch.cat(ys)
109

110

111
def train_mini_mcd(use_fddm=True, steps=400, seed=0):
112
    torch.manual_seed(seed)
113
    xs, ys = make_iq_short()
114
    n_tr = int(len(ys) * 0.8)
115
    x_tr, y_tr = xs[:n_tr], ys[:n_tr]
116
    x_te, y_te = xs[n_tr:], ys[n_tr:]
117

118
    model = MiniMCDformer(L=128, num_classes=4, use_fddm=use_fddm)
119
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
120
    for step in range(steps):
121
        idx = torch.randint(0, len(x_tr), (64,))
122
        opt.zero_grad()
123
        loss = F.cross_entropy(model(x_tr[idx]), y_tr[idx])
124
        loss.backward(); opt.step()
125

126
    model.eval()
127
    with torch.no_grad():
128
        acc = (model(x_te).argmax(1) == y_te).float().mean()
129
    return acc.item()
130

131

132
if __name__ == "__main__":
133
    acc_on = train_mini_mcd(use_fddm=True)
134
    acc_off = train_mini_mcd(use_fddm=False)
135
    print(f"低 SNR 4 类识别:  带 FDDM acc = {acc_on:.3f}    无 FDDM acc = {acc_off:.3f}")
136
    print("(随机猜测 = 0.250；FDDM 差距随噪声增大而扩大)")