### Python 项目：实现论文中提出的轻量级空间-光谱移位模块与多头 MambaOut (LS2SM-MHMambaOut) #### 项目概述 本项目旨在实现并复现论文 "Lightweight Spatial–Spectral Shift Module With Multihead MambaOut for Hyperspectral Image Classification" 中提出的 LS2SM-MHMambaOut 模型。该模型主要用于高光谱图像分类，通过引入轻量级的空间-光谱移位模块和多头 MambaOut 网络结构，能够在保证分类精度的同时减少模型参数和计算复杂度。 #### 目录结构 ``` ls2sm_mhmambaout/ │ ├── data/ │ ├── indian_pines.npy │ ├── longkou.npy │ └── ... (other datasets) │ ├── models/ │ ├── spcsm.py │ ├── s2sm.py │ ├── mhcnn.py │ └── ls2sm_mhmambaout.py │ ├── utils/ │ ├── data_loader.py │ ├── train.py │ └── eval.py │ ├── main.py └── README.md ``` #### 主要模块实现 ##### 1. Spectral Shift Module (SPCSM) `models/spcsm.py` ```python import torch import torch.nn as nn import torch.nn.functional as F class SPCSM(nn.Module): def __init__(self, in_channels, out_channels): super(SPCSM, self).__init__() self.fc = nn.Linear(in_channels, out_channels) def forward(self, x): # Circular shift operation on spectral dimension shifted_left = torch.roll(x, shifts=-1, dims=1) shifted_right = torch.roll(x, shifts=1, dims=1) # Aggregate shifted features aggregated = shifted_left + x + shifted_right # Fully connected layer for spectral mapping x_hat = self.fc(aggregated.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) return x_hat ``` ##### 2. Spatial-Spectral Shift Module (S2SM) `models/s2sm.py` ```python import torch import torch.nn as nn from .spcsm import SPCSM class S2SM(nn.Module): def __init__(self, in_channels, out_channels): super(S2SM, self).__init__() self.spcsm = SPCSM(in_channels, out_channels) self.spasm = SPASM() def forward(self, x): # Extract spectral and spatial features x_spc = self.spcsm(x) x_spa = self.spasm(x) # Adaptive fusion of spectral and spatial features spa_attention = torch.softmax(x_spa, dim=1) spc_attention = torch.softmax(x_spc, dim=1) fused = spa_attention * x_spa + spc_attention * x_spc return fused ``` ##### 3. Multihead Convolutional Block (MHCNN) `models/mhcnn.py` ```python import torch import torch.nn as nn class MHCNN(nn.Module): def __init__(self, in_channels, out_channels): super(MHCNN, self).__init__() self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) self.norm = nn.LayerNorm(out_channels) self.ffn = nn.Sequential( nn.Conv2d(out_channels, out_channels, kernel_size=1), nn.ReLU(), nn.Conv2d(out_channels, out_channels, kernel_size=1) ) def forward(self, x): x_res = self.conv(x) x_norm = self.norm(x_res.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) x_out = self.ffn(x_norm) + x_res return x_out ``` ##### 4. LS2SM-MHMambaOut 网络 `models/ls2sm_mhmambaout.py` ```python import torch import torch.nn as nn from .s2sm import S2SM from .mhcnn import MHCNN class LS2SM_MHMambaOut(nn.Module): def __init__(self, in_channels, num_classes): super(LS2SM_MHMambaOut, self).__init__() self.sc_block = nn.Conv2d(in_channels, 96, kernel_size=1) self.sslfem = nn.Sequential( S2SM(96, 96), nn.Conv2d(96, 96, kernel_size=3, padding=1), nn.ReLU() ) self.mhcnn_blocks = nn.ModuleList([MHCNN(96, 96) for _ in range(3)]) self.classifier = nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(96, num_classes) ) def forward(self, x): x = self.sc_block(x) x = self.sslfem(x) for block in self.mhcnn_blocks: x = block(x) x = self.classifier(x) return x ``` ##### 5. 数据加载和训练脚本 `utils/data_loader.py` ```python import numpy as np import torch from torch.utils.data import Dataset, DataLoader class HSIDataset(Dataset): def __init__(self, data_path, label_path, patch_size=15): self.data = np.load(data_path) self.labels = np.load(label_path) self.patch_size = patch_size def __len__(self): return len(self.labels) def __getitem__(self, idx): patch = self.data[idx] label = self.labels[idx] patch_tensor = torch.tensor(patch, dtype=torch.float32) label_tensor = torch.tensor(label, dtype=torch.long) return patch_tensor, label_tensor def get_data_loader(data_path, label_path, batch_size=32, patch_size=15): dataset = HSIDataset(data_path, label_path, patch_size) loader = DataLoader(dataset, batch_size=batch_size, shuffle=True) return loader ``` `utils/train.py` ```python import torch import torch.optim as optim from torch.utils.tensorboard import SummaryWriter from .data_loader import get_data_loader from models.ls2sm_mhmambaout import LS2SM_MHMambaOut def train_model(model, train_loader, val_loader, epochs=150, lr=0.001): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model.to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=lr) writer = SummaryWriter() for epoch in range(epochs): model.train() running_loss = 0.0 for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() avg_loss = running_loss / len(train_loader) writer.add_scalar('Training Loss', avg_loss, epoch) if epoch % 10 == 0: print(f'Epoch {epoch}/{epochs}, Loss: {avg_loss:.4f}') validate_model(model, val_loader, epoch, writer) writer.close() def validate_model(model, val_loader, epoch, writer): model.eval() correct = 0 total = 0 with torch.no_grad(): for inputs, labels in val_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) _, predicted = torch.max(outputs, 1) total += labels.size(0) correct += (predicted == labels).sum().item() accuracy = 100 * correct / total writer.add_scalar('Validation Accuracy', accuracy, epoch) print(f'Validation Accuracy: {accuracy:.2f}%') ``` `main.py` ```python import torch from utils.train import train_model from utils.data_loader import get_data_loader from models.ls2sm_mhmambaout import LS2SM_MHMambaOut if __name__ == '__main__': # Load data train_loader = get_data_loader('data/indian_pines.npy', 'data/indian_pines_labels.npy') val_loader = get_data_loader('data/indian_pines_val.npy', 'data/indian_pines_val_labels.npy') # Initialize model model = LS2SM_MHMambaOut(in_channels=200, num_classes=16) # Train model train_model(model, train_loader, val_loader) ``` #### 项目运行 1. **准备数据集**： - 将高光谱图像数据集（如 Indian Pines, WHU-Hi-LongKou 等）保存为 `.npy` 文件格式，并放置在 `data/` 文件夹下。 2. **安装依赖库**： ```bash pip install torch torchvision tensorboard numpy ``` 3. **运行主程序**： ```bash python main.py ``` #### 项目说明 - **README.md**：详细描述项目的背景、结构、功能及使用方法。 - **utils/eval.py**：用于评估模型性能的模块，可以扩展以支持更多评价指标。 - **其他模块**：可以根据需求进一步扩展和优化，例如添加超参数搜索、模型保存与加载等功能。 #### 知识点 1. **卷积神经网络（CNN）**： - 使用卷积层提取空间特征。 - 线性操作减少参数数量，提高效率。 2. **注意力机制（Attention Mechanism）**： - 通过自注意力机制融合全局特征。 - 提高模型对重要信息的关注度。 3. **高光谱图像处理（Hyperspectral Image Processing）**： - 高光谱图像具有丰富的空间-光谱信息。 - 提取这些信息可以提升分类效果。 #### 总结 此项目实现了论文中的 LS2SM-MHMambaOut 网络结构，并提供了完整的数据加载、模型训练和验证流程。通过复现论文方法，可以在实际应用中验证其有效性和性能表现。