用PyTorch复现Spectral-Spatial Attention Network：一个遥感图像分类的保姆级实战教程

用PyTorch复现Spectral-Spatial Attention Network一个遥感图像分类的保姆级实战教程高光谱遥感图像分类一直是计算机视觉领域极具挑战性的任务。传统的机器学习方法在处理这类数据时往往捉襟见肘而深度学习的出现为这一领域带来了革命性的突破。本文将带你从零开始用PyTorch实现一个结合了光谱和空间注意力机制的先进网络模型完整覆盖数据预处理、模型构建、训练优化到结果可视化的全流程。1. 环境准备与数据加载在开始构建模型前我们需要配置合适的开发环境并准备高光谱数据集。推荐使用Python 3.8和PyTorch 1.10版本这些组合能提供良好的兼容性和性能表现。首先安装必要的依赖库pip install torch torchvision numpy scikit-learn matplotlib scipy对于高光谱数据我们将使用经典的Pavia University数据集。这个数据集包含610×340像素的图像具有103个光谱波段涵盖9种不同的地表覆盖类型。以下是加载和预处理数据的完整代码import numpy as np import scipy.io as sio from sklearn.decomposition import PCA from sklearn.model_selection import train_test_split from torch.utils.data import Dataset, DataLoader class HSI_Dataset(Dataset): def __init__(self, data, labels, patch_size27, pca_components3): self.data data self.labels labels self.patch_size patch_size self.pca PCA(n_componentspca_components) # 数据标准化 self.data (self.data - np.mean(self.data)) / np.std(self.data) # PCA降维 self.data_pca self.pca.fit_transform(self.data.reshape(-1, data.shape[-1])) self.data_pca self.data_pca.reshape(data.shape[0], data.shape[1], pca_components) def __len__(self): return np.count_nonzero(self.labels) def __getitem__(self, idx): # 获取带标签的像素坐标 coords np.argwhere(self.labels 0) row, col coords[idx] label self.labels[row, col] - 1 # 类别从0开始 # 提取空间patch half self.patch_size // 2 patch np.pad(self.data_pca, ((half,half),(half,half),(0,0)), constant) spatial_patch patch[row:rowself.patch_size, col:colself.patch_size, :] # 提取光谱向量 spectral_vector self.data[row, col, :] return (torch.FloatTensor(spectral_vector), torch.FloatTensor(spatial_patch.transpose(2,0,1)), torch.LongTensor([label]))提示在实际应用中建议将数据集划分为训练集、验证集和测试集比例通常为6:2:2。对于类别不平衡问题可以采用过采样或加权损失函数等策略。2. 模型架构设计Spectral-Spatial Attention Network的核心在于同时捕捉光谱和空间两个维度的特征并通过注意力机制强化关键信息。我们将模型分解为三个主要组件光谱注意力分支、空间注意力分支和特征融合模块。2.1 光谱注意力分支光谱分支采用双向GRU结构处理连续光谱信息配合注意力机制突出重要波段import torch.nn as nn class SpectralAttention(nn.Module): def __init__(self, input_dim, hidden_dim, num_layers1): super().__init__() self.gru_fw nn.GRU(input_dim, hidden_dim, num_layers, batch_firstTrue) self.gru_bw nn.GRU(input_dim, hidden_dim, num_layers, batch_firstTrue) # 注意力机制 self.attention nn.Sequential( nn.Linear(2*hidden_dim, hidden_dim), nn.Tanh(), nn.Linear(hidden_dim, 1), nn.Softmax(dim1) ) def forward(self, x): # 双向GRU处理 out_fw, _ self.gru_fw(x.unsqueeze(1)) # (batch, seq_len, hidden) out_bw, _ self.gru_bw(torch.flip(x.unsqueeze(1), [1])) out_bw torch.flip(out_bw, [1]) # 拼接双向输出 combined torch.cat([out_fw, out_bw], dim-1).squeeze(1) # (batch, 2*hidden) # 计算注意力权重 attn_weights self.attention(combined) attended (attn_weights * combined).sum(dim1) return attended, attn_weights.squeeze()2.2 空间注意力分支空间分支采用CNN架构处理局部邻域信息通过空间注意力强化关键区域class SpatialAttention(nn.Module): def __init__(self, in_channels): super().__init__() self.conv1 nn.Conv2d(in_channels, 32, kernel_size3, padding1) self.conv2 nn.Conv2d(32, 64, kernel_size3, padding1) self.pool nn.MaxPool2d(2, 2) # 空间注意力 self.attn_conv nn.Conv2d(64, 1, kernel_size1) self.sigmoid nn.Sigmoid() def forward(self, x): # 基础特征提取 x F.relu(self.conv1(x)) x self.pool(x) x F.relu(self.conv2(x)) # 空间注意力 attn self.sigmoid(self.attn_conv(x)) attended x * attn # 全局平均池化 out F.adaptive_avg_pool2d(attended, (1,1)).view(x.size(0), -1) return out, attn.squeeze()2.3 特征融合与分类将两个分支的特征进行融合后通过全连接层进行分类class SSANet(nn.Module): def __init__(self, spectral_dim, spatial_channels, num_classes): super().__init__() self.spectral_branch SpectralAttention(spectral_dim, 64) self.spatial_branch SpatialAttention(spatial_channels) # 融合分类 self.fc nn.Sequential( nn.Linear(128 64, 256), nn.ReLU(), nn.Dropout(0.5), nn.Linear(256, num_classes) ) def forward(self, spectral, spatial): # 光谱分支 spec_feat, spec_attn self.spectral_branch(spectral) # 空间分支 spat_feat, spat_attn self.spatial_branch(spatial) # 特征融合 combined torch.cat([spec_feat, spat_feat], dim1) logits self.fc(combined) return logits, spec_attn, spat_attn3. 模型训练与优化构建好模型架构后我们需要设计合适的训练流程和优化策略。高光谱数据通常样本有限因此需要特别注意防止过拟合。3.1 损失函数与评估指标对于多分类问题交叉熵损失是标准选择。同时监控准确率和Kappa系数def train_model(model, dataloaders, criterion, optimizer, num_epochs100): best_acc 0.0 for epoch in range(num_epochs): for phase in [train, val]: if phase train: model.train() else: model.eval() running_loss 0.0 running_corrects 0 for spectral, spatial, labels in dataloaders[phase]: spectral spectral.to(device) spatial spatial.to(device) labels labels.to(device).squeeze() optimizer.zero_grad() with torch.set_grad_enabled(phase train): outputs, _, _ model(spectral, spatial) loss criterion(outputs, labels) if phase train: loss.backward() optimizer.step() _, preds torch.max(outputs, 1) running_loss loss.item() * spectral.size(0) running_corrects torch.sum(preds labels.data) epoch_loss running_loss / len(dataloaders[phase].dataset) epoch_acc running_corrects.double() / len(dataloaders[phase].dataset) print(f{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}) # 保存最佳模型 if phase val and epoch_acc best_acc: best_acc epoch_acc torch.save(model.state_dict(), best_model.pth) return model3.2 学习率调度与正则化采用余弦退火学习率调度和标签平滑技术提升模型泛化能力from torch.optim.lr_scheduler import CosineAnnealingLR from torch.nn.functional import cross_entropy class LabelSmoothingCrossEntropy(nn.Module): def __init__(self, epsilon0.1): super().__init__() self.epsilon epsilon def forward(self, logits, targets): num_classes logits.size(-1) log_probs F.log_softmax(logits, dim-1) targets torch.zeros_like(log_probs).scatter_(1, targets.unsqueeze(1), 1) targets (1 - self.epsilon) * targets self.epsilon / num_classes loss (-targets * log_probs).sum(dim1).mean() return loss # 初始化 model SSANet(spectral_dim103, spatial_channels3, num_classes9).to(device) criterion LabelSmoothingCrossEntropy() optimizer torch.optim.Adam(model.parameters(), lr0.001, weight_decay1e-4) scheduler CosineAnnealingLR(optimizer, T_max10, eta_min1e-5)4. 结果分析与可视化训练完成后我们需要评估模型性能并理解其决策过程。注意力机制的一个优势就是提供了可解释性。4.1 分类性能评估在测试集上计算混淆矩阵和各类别指标from sklearn.metrics import confusion_matrix, classification_report def evaluate_model(model, test_loader): model.eval() all_preds [] all_labels [] with torch.no_grad(): for spectral, spatial, labels in test_loader: spectral spectral.to(device) spatial spatial.to(device) labels labels.to(device).squeeze() outputs, _, _ model(spectral, spatial) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) # 计算评估指标 cm confusion_matrix(all_labels, all_preds) report classification_report(all_labels, all_preds, target_namesclass_names) print(Confusion Matrix:) print(cm) print(\nClassification Report:) print(report) return cm, report4.2 注意力可视化绘制光谱和空间注意力图理解模型关注的重点import matplotlib.pyplot as plt def visualize_attention(model, sample): model.eval() spectral, spatial, label sample spectral spectral.unsqueeze(0).to(device) spatial spatial.unsqueeze(0).to(device) with torch.no_grad(): _, spec_attn, spat_attn model(spectral, spatial) # 光谱注意力 plt.figure(figsize(12,4)) plt.subplot(1,2,1) plt.plot(spectral.squeeze().cpu().numpy(), label光谱曲线) plt.plot(spec_attn.squeeze().cpu().numpy(), label注意力权重) plt.title(光谱注意力) plt.legend() # 空间注意力 plt.subplot(1,2,2) plt.imshow(spat_attn.squeeze().cpu().numpy(), cmaphot) plt.title(空间注意力热图) plt.colorbar() plt.show()5. 实战技巧与常见问题在实际复现过程中有几个关键点需要特别注意数据增强高光谱数据有限可以通过旋转、翻转等方式增加样本多样性梯度裁剪RNN容易出现梯度爆炸设置nn.utils.clip_grad_norm_控制梯度范围混合精度训练使用torch.cuda.amp加速训练并减少显存占用早停机制监控验证集损失当连续若干轮不下降时停止训练常见问题及解决方案问题现象可能原因解决方案训练损失不下降学习率设置不当尝试不同学习率或使用学习率查找器验证集准确率波动大批次大小不合适增大批次大小或使用梯度累积测试集性能差过拟合增加Dropout比例或使用更多正则化注意力权重集中模型退化检查初始化方式添加残差连接在Pavia University数据集上的实验表明完整的SSANet模型能够达到约98.2%的总体准确率相比单独使用CNN或RNN有显著提升。光谱注意力机制成功识别出对分类贡献最大的波段区域而空间注意力则有效聚焦于目标物体的中心区域。