当前位置：首页 > news >正文

Transformer也能玩转高光谱图像分类？手把手教你复现SpectralFormer（附代码）

news 2026/6/13 17:12:27

Transformer在高光谱图像分类中的实战应用：SpectralFormer完整复现指南

高光谱图像分类一直是遥感领域的重要研究方向，而Transformer架构的引入为这一领域带来了全新的可能性。本文将带您深入探索SpectralFormer这一创新模型，从理论到实践，手把手教您完成整个复现流程。

1. 环境准备与数据预处理

复现SpectralFormer的第一步是搭建合适的开发环境。我们推荐使用Python 3.8+和PyTorch 1.9+作为基础框架，同时需要安装一些必要的依赖库：

pip install torch torchvision numpy scipy scikit-learn matplotlib tqdm

对于高光谱数据集，我们主要使用三个经典基准数据集：Indian Pines、Pavia University和Houston 2013。这些数据集可以从公开资源获取，下载后需要进行以下预处理步骤：

数据标准化：对每个光谱波段进行Z-score标准化处理
数据划分：按照论文中的标准划分训练集和测试集
数据增强（可选）：对训练数据应用随机旋转、翻转等空间变换

import numpy as np from sklearn.preprocessing import StandardScaler def preprocess_data(data): # 数据标准化 original_shape = data.shape data_2d = data.reshape(-1, original_shape[-1]) scaler = StandardScaler() data_normalized = scaler.fit_transform(data_2d) return data_normalized.reshape(original_shape)

2. SpectralFormer架构深度解析

SpectralFormer的核心创新在于其独特的GroupWise频谱嵌入和跨层自适应融合机制。让我们深入剖析这两个关键模块的实现细节。

2.1 GroupWise频谱嵌入实现

传统的Transformer处理高光谱数据时，通常将每个波段视为独立的token。而SpectralFormer则采用GroupWise方式，将相邻波段组合成组进行处理：

import torch import torch.nn as nn class GroupWiseEmbedding(nn.Module): def __init__(self, in_channels, embed_dim, group_size=3): super().__init__() self.group_size = group_size self.projection = nn.Linear(in_channels * group_size, embed_dim) def forward(self, x): # x shape: [batch, bands, channels] b, n, c = x.shape # 分组处理 x = x.unfold(1, self.group_size, 1) # [b, n-g+1, c, g] x = x.permute(0,1,3,2).contiguous() # [b, n-g+1, g, c] x = x.view(b, -1, self.group_size * c) # [b, n-g+1, g*c] # 投影到嵌入空间 return self.projection(x)

2.2 跨层自适应融合模块

跨层自适应融合(CAF)是SpectralFormer的另一大创新，它通过可学习的权重参数自适应地融合不同层的特征：

class CrossLayerFusion(nn.Module): def __init__(self, dim): super().__init__() self.fusion_weights = nn.Parameter(torch.randn(2, dim)) self.norm = nn.LayerNorm(dim) def forward(self, prev_features, current_features): # prev_features: 前几层的特征 # current_features: 当前层特征 fused = torch.stack([prev_features, current_features], dim=-1) weights = torch.softmax(self.fusion_weights, dim=0) fused_features = torch.matmul(fused, weights) return self.norm(fused_features)

3. 完整模型搭建与训练策略

基于上述核心模块，我们可以构建完整的SpectralFormer模型。以下是模型的主要架构：

class SpectralFormer(nn.Module): def __init__(self, num_classes, num_bands, embed_dim=64, depth=5, num_heads=4, group_size=3): super().__init__() # 频谱嵌入层 self.embedding = GroupWiseEmbedding(1, embed_dim, group_size) # Transformer编码器层 encoder_layer = nn.TransformerEncoderLayer( d_model=embed_dim, nhead=num_heads) self.transformer = nn.TransformerEncoder(encoder_layer, depth) # 分类头 self.classifier = nn.Linear(embed_dim, num_classes) # 跨层融合模块 self.cafs = nn.ModuleList([ CrossLayerFusion(embed_dim) for _ in range(depth//2)]) def forward(self, x): # x shape: [batch, bands] x = x.unsqueeze(-1) # [b, bands, 1] x = self.embedding(x) # [b, n, embed_dim] # 保存中间层特征用于跨层融合 features = [] for i, layer in enumerate(self.transformer.layers): x = layer(x) if i % 2 == 1 and i > 0: # 每隔两层应用一次CAF x = self.cafs[i//2 - 1](features[-1], x) features.append(x) # 全局平均池化后分类 x = x.mean(dim=1) return self.classifier(x)

3.1 训练配置与优化策略

为了获得最佳性能，我们需要精心配置训练参数：

优化器：AdamW优化器，初始学习率5e-4
学习率调度：余弦退火策略
正则化：权重衰减5e-3，Dropout率0.1
批次大小：64（根据GPU显存调整）

from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR model = SpectralFormer(num_classes=16, num_bands=200) optimizer = AdamW(model.parameters(), lr=5e-4, weight_decay=5e-3) scheduler = CosineAnnealingLR(optimizer, T_max=1000) # 训练循环示例 for epoch in range(1000): model.train() for x, y in train_loader: optimizer.zero_grad() outputs = model(x) loss = criterion(outputs, y) loss.backward() optimizer.step() scheduler.step()

4. 实战技巧与性能优化

在实际复现过程中，可能会遇到各种挑战。以下是几个关键问题的解决方案：

4.1 显存不足问题处理

高光谱数据通常需要较大显存，特别是处理空间-光谱立方体时。可以采用以下策略：

梯度累积：小批次训练，多次累积后更新
混合精度训练：使用AMP自动混合精度
数据分块：将大图像分割为小块处理

from torch.cuda.amp import GradScaler, autocast scaler = GradScaler() for x, y in train_loader: optimizer.zero_grad() with autocast(): outputs = model(x) loss = criterion(outputs, y) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

4.2 模型收敛问题

SpectralFormer训练初期可能出现不稳定现象，可以通过以下方法改善：

学习率预热：前10个epoch线性增加学习率
标签平滑：减轻过拟合
梯度裁剪：防止梯度爆炸

# 标签平滑实现 class LabelSmoothingLoss(nn.Module): def __init__(self, smoothing=0.1): super().__init__() self.smoothing = smoothing def forward(self, logits, targets): n_classes = logits.size(-1) log_preds = F.log_softmax(logits, dim=-1) loss = -log_preds.mean() nll = F.nll_loss(log_preds, targets) return (1 - self.smoothing) * nll + self.smoothing * loss

4.3 评估指标实现

高光谱分类常用评估指标包括总体精度(OA)、平均精度(AA)和Kappa系数：

from sklearn.metrics import confusion_matrix, cohen_kappa_score def evaluate(model, loader): model.eval() all_preds, all_labels = [], [] with torch.no_grad(): for x, y in loader: outputs = model(x) preds = outputs.argmax(dim=1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(y.cpu().numpy()) cm = confusion_matrix(all_labels, all_preds) oa = np.sum(np.diag(cm)) / np.sum(cm) aa = np.mean(np.diag(cm) / np.sum(cm, axis=1)) kappa = cohen_kappa_score(all_labels, all_preds) return oa, aa, kappa

5. 进阶应用与扩展思考

掌握了基础复现方法后，我们可以进一步探索SpectralFormer的潜力：

5.1 空间-光谱联合建模

原始SpectralFormer主要处理光谱信息，我们可以扩展其处理空间信息的能力：

空间注意力机制：在Transformer中加入空间注意力头
多尺度特征融合：结合不同尺度的空间特征
三维卷积预处理：先用3D CNN提取空间-光谱特征

class SpatialSpectralFormer(nn.Module): def __init__(self, num_classes, patch_size=7): super().__init__() self.patch_embed = nn.Conv2d(1, 64, kernel_size=patch_size, stride=patch_size) self.spectral_embed = GroupWiseEmbedding(64, 64) self.transformer = nn.TransformerEncoder(...) def forward(self, x): # x: [b, c, h, w] patches = self.patch_embed(x) # [b, e, h', w'] b, e, h, w = patches.shape patches = patches.permute(0,2,3,1).reshape(b, h*w, e) spectral_emb = self.spectral_embed(patches) return self.transformer(spectral_emb)

5.2 轻量化设计

针对实际应用中的效率需求，可以考虑以下优化方向：

知识蒸馏：用大模型训练小模型
结构剪枝：移除不重要的注意力头或层
量化感知训练：准备模型用于8位整数量化

# 知识蒸馏示例 def distillation_loss(student_logits, teacher_logits, labels, temp=2.0, alpha=0.5): soft_teacher = F.softmax(teacher_logits/temp, dim=1) soft_student = F.log_softmax(student_logits/temp, dim=1) kl_div = F.kl_div(soft_student, soft_teacher, reduction='batchmean') ce_loss = F.cross_entropy(student_logits, labels) return alpha * kl_div + (1 - alpha) * ce_loss

5.3 自监督预训练

针对高光谱数据标注成本高的问题，可以探索自监督预训练策略：

波段预测：随机mask部分波段进行预测
对比学习：构建正负样本对进行对比
拼图重建：打乱空间位置后重建

class MaskedBandPrediction(nn.Module): def __init__(self, encoder): super().__init__() self.encoder = encoder self.pred_head = nn.Linear(64, 1) # 预测被mask的波段值 def forward(self, x, mask_ratio=0.2): # 随机mask部分波段 b, n, c = x.shape mask = torch.rand(b, n) > mask_ratio masked_x = x * mask.unsqueeze(-1) features = self.encoder(masked_x) preds = self.pred_head(features) return preds, x[~mask] # 返回预测值和真实值

6. 实际应用中的挑战与解决方案

在将SpectralFormer应用于实际项目时，可能会遇到一些特有的挑战：

6.1 小样本学习

高光谱分类常面临标注数据稀缺的问题。可以采用以下策略：

数据增强：专门设计的光谱变换（如高斯噪声、波段dropout）
迁移学习：在大型数据集上预训练，在小数据集上微调
半监督学习：利用未标注数据提升性能

class SpectralAugmentation: def __init__(self, noise_std=0.1, dropout_prob=0.1): self.noise_std = noise_std self.dropout_prob = dropout_prob def __call__(self, x): # 添加高斯噪声 if self.noise_std > 0: x = x + torch.randn_like(x) * self.noise_std # 随机丢弃部分波段 if self.dropout_prob > 0: mask = torch.rand(x.shape[1]) > self.dropout_prob x = x * mask.float() return x

6.2 类别不平衡处理

高光谱数据中不同类别样本数量可能差异很大。解决方法包括：

加权损失函数：根据类别频率调整损失权重
过采样/欠采样：平衡各类别样本数量
焦点损失：降低易分类样本的权重

# 加权交叉熵损失 def weighted_cross_entropy(logits, labels, class_weights): log_probs = F.log_softmax(logits, dim=-1) weights = class_weights[labels] return -(weights * log_probs[range(len(labels)), labels]).mean() # 计算类别权重 def compute_class_weights(labels): class_counts = torch.bincount(labels) return 1.0 / (class_counts.float() / class_counts.sum())

6.3 跨场景泛化

在不同地点采集的高光谱数据分布可能差异很大。提升模型泛化能力的方法：

领域自适应：使用对抗训练对齐特征分布
光谱归一化：减少传感器差异影响
元学习：学习快速适应新场景的能力

# 领域鉴别器 class DomainDiscriminator(nn.Module): def __init__(self, input_dim): super().__init__() self.net = nn.Sequential( nn.Linear(input_dim, 64), nn.ReLU(), nn.Linear(64, 1)) def forward(self, features): return self.net(features.detach()) # 对抗训练损失 def adversarial_loss(source_feat, target_feat, discriminator): source_dom = discriminator(source_feat) target_dom = discriminator(target_feat) loss = F.binary_cross_entropy_with_logits( torch.cat([source_dom, target_dom]), torch.cat([torch.ones_like(source_dom), torch.zeros_like(target_dom)])) return loss