SSFTT网络：高光谱图像分类的3D卷积与Transformer融合

1. SSFTT网络架构解析与代码实现

SSFTT(Spectral-Spatial Feature Tokenization Transformer)是一种专门用于高光谱图像分类的深度学习模型。这个架构巧妙地将3D卷积、2D卷积与Transformer相结合，能够同时捕捉光谱和空间特征。下面我将从代码实现角度详细解析这个网络。

1.1 模型整体架构

SSFTT网络的核心思想可以概括为四个关键步骤：

1. 使用3D卷积提取初步的光谱-空间特征
2. 通过2D卷积进一步提炼空间特征
3. 将特征图转换为token序列
4. 使用Transformer进行全局关系建模

在代码中，这对应SSFTTnet.py文件中的SSFTTnet类。让我们先看模型初始化的关键参数：

python复制class SSFTTnet(nn.Module):
    def __init__(self, band, classes, num_tokens=4):
        super(SSFTTnet, self).__init__()
        # 光谱-空间特征提取
        self.conv3d_1 = nn.Sequential(
            nn.Conv3d(1, 8, (7, 3, 3), padding=(0, 1, 1)),
            nn.BatchNorm3d(8),
            nn.ReLU()
        )
        # 空间特征提取
        self.conv2d = nn.Sequential(
            nn.Conv2d(224, 64, (3, 3)),
            nn.BatchNorm2d(64),
            nn.ReLU()
        )
        # Tokenizer
        self.token_wA = nn.Parameter(torch.empty(64, 64))
        # Transformer编码器
        self.transformer = Transformer(dim=64, depth=2, heads=8, dim_head=64, mlp_dim=128)
        # 分类头
        self.cls_token = nn.Parameter(torch.randn(1, 1, 64))
        self.mlp_head = nn.Linear(64, classes)

1.2 输入输出维度变化

理解维度变化是掌握模型的关键。让我们跟踪一个典型输入(64,1,30,13,13)在各层的变换：

1. 初始输入：(batch=64, channel=1, spectral=30, height=13, width=13)
2. Conv3D后：(64,8,24,11,11)
   • 3D卷积核(7,3,3)在光谱维度压缩30→24
3. 重排reshape：(64,192,11,11)
   • 将8×24=192个特征图展平
4. Conv2D后：(64,64,9,9)
   • 空间维度从11×11→9×9
5. Tokenizer后：(64,4,64)
   • 将81个空间位置(9×9)聚合为4个token
6. 加CLS后：(64,5,64)
   • 添加1个分类token
7. Transformer后：(64,5,64)
   • 保持维度不变但内部特征已交互
8. 分类头输出：(64,classes)
   • 只取CLS token进行分类

提示：实际调试时可以在forward函数中添加print语句，实时查看各层维度变化。这是理解复杂模型最有效的方法之一。

2. 核心模块实现细节

2.1 光谱-空间特征提取

3D卷积的设计是SSFTT的第一个关键点：

python复制nn.Conv3d(1, 8, (7, 3, 3), padding=(0, 1, 1))

这个卷积核在三个维度上的含义：

• 光谱维度：7 - 跨7个波段提取光谱特征
• 空间维度：3×3 - 标准空间卷积
• padding=(0,1,1)表示只在空间维度填充，保持空间尺寸

为什么需要先3D再2D？

1. 3D卷积能同时捕捉邻近波段和空间邻域的相关性
2. 但3D计算量巨大，不宜过深
3. 2D卷积专注于空间特征提取，计算更高效

2.2 Tokenizer实现

Tokenizer是将特征图转换为token的关键模块，对应论文中的公式(3)：

python复制# 代码实现
wa = rearrange(self.token_wA, 'd h -> h d')  # 64,64
a = torch.einsum('bnd,h->bnh', x, wa)       # 64,81,64
a = a.softmax(dim=1)                        # 64,81,64
tokens = torch.einsum('bnh,bnd->bhd', a, x) # 64,4,64

数学表达式：
T = softmax(XWₐ)ᵀX

其中：

• X ∈ ℝ^{B×N×D} 是输入特征(B:batch, N:空间位置, D:特征维度)
• Wₐ ∈ ℝ^{D×H} 是可学习参数
• softmax沿N维度归一化
• 输出T ∈ ℝ^{B×H×D} 是生成的token

选择softmax的原因：

1. 产生归一化的注意力权重
2. 使token对输入特征的聚合更平滑
3. 可微分，适合端到端训练

2.3 Transformer编码器

Transformer部分采用标准实现，但针对高光谱数据做了调整：

python复制class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Sequential(
                PreNorm(dim, Attention(dim, heads, dim_head)),
                PreNorm(dim, FeedForward(dim, mlp_dim))
            ) for _ in range(depth)
        ])

关键参数选择：

• dim=64：与token维度一致
• depth=2：较浅的层数防止过拟合
• heads=8：多头注意力提升表征能力
• mlp_dim=128：隐藏层扩展比例2倍

3. 训练流程与数据准备

3.1 数据预处理流程

IP_train.py中的完整流程：

1. 数据读取：从.mat文件加载高光谱图像和标签
2. PCA降维：保留前30个主成分(可配置)
3. Patch提取：以每个像素为中心取13×13邻域
4. 数据集划分：按比例分为训练/验证/测试集
5. 数据增强：随机旋转、翻转增加多样性

关键代码片段：

python复制# PCA降维
pca = PCA(n_components=30)
data = pca.fit_transform(data.reshape(-1, bands)).reshape(h, w, 30)

# Patch提取
patches = []
for i in range(h):
    for j in range(w):
        patch = data[max(0,i-6):min(h,i+7), max(0,j-6):min(w,j+7)]
        patch = np.pad(patch, ((6-i, i+7-h), (6-j, j+7-w), (0,0)), 'constant')
        patches.append(patch)

3.2 训练配置

训练使用标准交叉熵损失和Adam优化器：

python复制criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=0.0001)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.9)

关键训练技巧：

1. 学习率衰减：每20epoch乘以0.9
2. 早停机制：验证集精度不再提升时停止
3. 梯度裁剪：防止梯度爆炸
4. 混合精度训练：加速训练过程

4. 结果可视化与分析

4.1 分类结果可视化

get_cls_map.py将预测结果转换为分类图：

python复制def generate_cls_map(model, data_loader, device):
    model.eval()
    cls_map = np.zeros((h, w))
    with torch.no_grad():
        for batch in data_loader:
            inputs, positions = batch
            outputs = model(inputs.to(device))
            preds = outputs.argmax(dim=1)
            for i, pos in enumerate(positions):
                cls_map[pos[0], pos[1]] = preds[i].item()
    return cls_map

可视化技巧：

1. 使用matplotlib的imshow显示分类图
2. 为不同类别分配鲜明颜色
3. 叠加原始图像作为背景提高可读性

4.2 性能评估指标

常用评估指标实现：

python复制def evaluate(model, data_loader, device):
    model.eval()
    total_correct = 0
    total_samples = 0
    conf_matrix = np.zeros((classes, classes))
    
    with torch.no_grad():
        for inputs, labels in data_loader:
            outputs = model(inputs.to(device))
            preds = outputs.argmax(dim=1)
            total_correct += (preds == labels.to(device)).sum().item()
            total_samples += labels.size(0)
            for t, p in zip(labels, preds):
                conf_matrix[t, p] += 1
                
    accuracy = total_correct / total_samples
    return accuracy, conf_matrix

5. 关键问题与解决方案

5.1 常见训练问题

1. 梯度消失/爆炸

   • 解决方案：使用LayerNorm、梯度裁剪
   • 代码实现：

     python复制torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
     

2. 过拟合

   • 解决方案：增加权重衰减、Dropout层
   • 代码位置：

     python复制self.dropout = nn.Dropout(0.5)  # 在Transformer中添加
     

3. 类别不平衡

   • 解决方案：加权交叉熵损失
   • 实现方式：

     python复制weights = torch.tensor([...], device=device)  # 每个类别的权重
     criterion = nn.CrossEntropyLoss(weight=weights)
     

5.2 模型调试技巧

1. 特征可视化

   • 使用PCA或t-SNE降维后绘制特征分布
   • 检查各层特征是否具有判别性

2. 学习率测试

   • 实施学习率范围测试：

   python复制lr_finder = LRFinder(model, optimizer, criterion)
   lr_finder.range_test(train_loader, end_lr=1, num_iter=100)
   lr_finder.plot()
   

3. 激活统计

   • 监控各层激活值的均值和方差：

   python复制print(f'Layer {name}: mean={activation.mean().item():.4f}, std={activation.std().item():.4f}')
   

6. 扩展与优化方向

6.1 计算效率优化

1. 混合精度训练

   python复制scaler = torch.cuda.amp.GradScaler()
   with torch.cuda.amp.autocast():
       outputs = model(inputs)
       loss = criterion(outputs, labels)
   scaler.scale(loss).backward()
   scaler.step(optimizer)
   scaler.update()
   

2. 模型轻量化

   • 使用深度可分离卷积
   • 减少token数量
   • 知识蒸馏

6.2 性能提升方向

1. 多尺度特征融合

   • 在不同层级提取patch
   • 跨尺度特征交互

2. 时序信息利用

   • 对多时相数据建模
   • 3D卷积扩展为(时间,空间,光谱)

3. 不确定性建模

   • 输出分类置信度
   • 对低置信度区域特殊处理

在实际项目中，我发现模型的性能很大程度上取决于数据预处理的质量。特别是PCA降维的维度选择和patch大小的设置需要反复实验。另一个关键点是初始学习率的设置，使用学习率范围测试可以显著减少调参时间。