PyTorch实现高光谱图像分类:2D_CNN从入门到实战
发布时间:2026/10/2 14:22:12来源:尧图网络
在深度学习做遥感分类的圈子里高光谱图像一直是个特别的存在。大家一听到“高光谱”三个字第一反应往往是数据维度高、模型要够复杂。我做这个项目最深的体会恰恰相反真正让我稳定拿到好结果的不是花哨的三维卷积而是一个结构非常朴素的2D_CNN。这篇就用手把手的方式讲讲怎么用PyTorch把高光谱图像分类这件事跑通从环境搭建、数据预处理到网络训练和评估全程可复现。如果你手头正拿着Indian Pines这样的经典高光谱数据集想快速出一个不错的分类结果或者刚接触高光谱图像分类、想搞清楚2D_CNN网络到底怎么吃三维数据这篇文章应该能帮你省掉不少弯路。1. 先搞清楚:高光谱图像分类到底在分什么1.1 从一张“三维数据”说起我们平时处理的自然图像基本是三个通道红、绿、蓝。高光谱图像完全不同它把光谱维拉得非常长。以最常用的Indian Pines数据集为例它的原始数据是一个145x145x220的三维数组通常去掉噪声和吸水波段后还剩200个波段。也就是说每个像素不再只有RGB三个值而是拥有一条200维的光谱曲线。如果把整个数据集想象成一个立方体横纵两个方向是空间位置竖着的是光谱反射率那每个地物类别在这条光谱曲线上都有自己独特的“指纹”。高光谱图像分类的任务就是根据每个像素的光谱曲线——以及它周围邻域的空间纹理——判断这个像素属于哪一类地物。Indian Pines里一共有16类比如玉米、大豆、小麦、森林、建筑等。这里有一个关键认知虽然数据是三维的但我们完全可以把它转化成“多通道的二维图像”来喂给卷积网络。这也是2D_CNN能在这个任务上站稳脚跟的根本原因。1.2 2D CNN的输入构造逻辑要让2D_CNN处理高光谱数据核心思路是把光谱维折叠成图像的通道维。具体做法是对每个像素取它周围一个固定大小的空间邻域窗口比如11x11。这个窗口里有121个像素每个像素有200个光谱值。我们把光谱维压缩通常用主成分分析PCA到30个主成分于是每个像素就相当于一张11x11大小、30个通道的“图像”。这个图像块就是2D_CNN的一个输入样本。所以整个数据集的处理流程是遍历所有有标签的像素以每个像素为中心裁出邻域patch对光谱维做降维得到(N, 30, 11, 11)的张量这个张量就可以直接输入一个标准的2D_CNN网络你可能会疑问光谱维压缩到30个通道会不会丢信息我做实验的经验是200个波段之间往往高度相关PCA保留30个主成分通常已经能覆盖95%以上的方差。压缩掉的更多是噪声不是有效信息。1.3 和3D CNN的真实差距很多人看到高光谱数据是三维立方体第一直觉是用3D卷积。我不否认3D CNN在理论上的表达能力更强它可以同时建模空间和光谱维的联合特征。但实际跑下来你会发现3D网络有几个难以回避的问题第一参数量暴涨。3D卷积核引入了额外的时间/光谱维同样的卷积核尺寸和深度下参数量成倍增加。Indian Pines训练样本总共才一万出头3D网络太容易过拟合。第二训练不稳定。我在实验里尝试过一个中等规模的3D CNN同样的数据划分和训练策略下收敛速度明显慢于2D版本而且最终精度在多次随机种子下波动更大。第三收益有限。遥感领域有个经验性的结论在样本量不足够大的情况下2D CNN配合PCA或者简单的空间-光谱特征提取往往能达到3D网络八九成的效果但训练成本只有三分之一。所以标题里的“超简单”并非噱头而是选型上的刻意为之。先把2D_CNN这条路走通拿到一个可靠的基线结果后续再往3D、混合注意力方向发展心里才有底。2. 环境准备与数据预处理把三维数据变成二维卷积能吃的样子2.1 PyTorch环境搭建Anaconda GPU这个项目的技术栈很明确Python 3.9 PyTorch配合Anaconda管理环境。为什么用Anaconda因为它能干净地隔离不同项目的依赖你在这个环境里装什么都不会影响系统Python踩坏了直接删掉重建十分钟就能恢复。创建环境的命令就这几条conda create -n hsi python3.9 conda activate hsi安装PyTorch这一步是最容易出问题的。版本号必须和你机器的CUDA驱动对得上否则即使装好了也会出现torch.cuda.is_available()返回False的尴尬情况。我的建议是直接去PyTorch官网挑对应的安装命令别自己去拼版本号。以目前常用的CUDA 11.8为例pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118如果你用的是NVIDIA显卡装完先做一次体检import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))三条输出都正常环境就稳了。如果你在Windows上训练时总遇到驱动兼容问题可以试试WSL2很多和我交流过的同学反馈WSL2下训练PyTorch的稳定性比原生Windows好不少。另外如果你用的是AMD显卡比如Radeon 7900XTX则需要走ROCm路线安装命令和CUDA版本不一样记得去PyTorch官网选对应平台。2.2 读取Indian Pines的.mat数据Indian Pines的原始文件是一个.mat格式的MATLAB文件。读取它不需要MATLAB用SciPy就行非常轻量。from scipy.io import loadmat import numpy as np # 替换成你的数据路径 data loadmat(data/Indian_pines_corrected.mat)[indian_pines_corrected] gt loadmat(data/Indian_pines_gt.mat)[indian_pines_gt] print(data.shape) # (145, 145, 200) print(gt.shape) # (145, 145)数据是200个波段的版本已经去除了吸水带噪声波段。ground truth里0代表背景像素1到16是16个地物类别。后续处理标签时要记住背景像素不参与训练也不参与评估。2.3 空间Patch提取的核心步骤这是整个预处理里最绕的一步我给你讲透。模型的每个输入样本不是单独一个像素的光谱曲线而是以该像素为中心、周围一圈邻域组成的“图像块”。假设窗口大小是11那每个样本的形状就是(11, 11, 200)。问题来了图像边缘的像素怎么办比如坐标(0,0)这个像素它的左边和上边都没有邻居。常见的做法是做边界填充。我在代码里用reflect填充模式效果比零填充好不少因为反射填充保持了图像边缘的结构连续性。你可以在代码里对比一下这两种模式对边缘类别精度的差异。处理流程拆成四步给原始数据加一圈paddingpadding大小等于窗口半径。11x11窗口的半径是5所以给数据四周各加5行/列。遍历ground truth中所有非0的像素位置。以每个位置为中心从padding后的数据里切出11x11x200的块。把块的标签和坐标记下来。核心代码长这样from sklearn.decomposition import PCA def pad_with_reflect(X, margin): return np.pad(X, ((margin, margin), (margin, margin), (0, 0)), modereflect) def extract_patches(X, gt, patch_size11): margin patch_size // 2 X_padded pad_with_reflect(X, margin) row, col np.nonzero(gt) patches [] labels [] positions [] for i, j in zip(row, col): patch X_padded[i:ipatch_size, j:jpatch_size, :] patches.append(patch) labels.append(gt[i, j]) positions.append((i, j)) return np.array(patches), np.array(labels), positions这一步跑完后patches的形状是(N, 11, 11, 200)其中N大约等于10249这是Indian Pines全部有标签的像素数。2.4 归一化与PCA顺序错了会出事很多入门教程把归一化和PCA直接放在划分训练测试集之前代码跑通了结果也好看但严谨地说这里存在数据泄漏风险。正确的做法是先划分训练集和测试集再用训练集的统计量去归一化测试集。这个细节后面我会在踩坑章节专门展开。PCA降维的逻辑也值得说两句。印度松林这200个相邻波段相关性很高直接全部喂进去模型要学的东西太多参数也压不住。PCA能把200维压缩成30维这本质上是一个无损或者说近无损的信息提纯过程。你用sklearn.decomposition.PCA就行但注意要把三维patch展平成二维再做PCA。完整预处理流水线from sklearn.decomposition import PCA from sklearn.preprocessing import StandardScaler # 先划分训练测试索引 # 这里用固定种子保证可复现详情见第4章 train_idx, test_idx train_test_split_indices(labels, seed3407) # 第一步在训练集上fit归一化器 train_flat patches[train_idx].reshape(train_idx.sum(), -1) scaler StandardScaler().fit(train_flat) # 第二步归一化全体数据统计量来自训练集 patches_flat patches.reshape(len(patches), -1) patches_norm scaler.transform(patches_flat) # 第三步在训练集上fit PCA pca PCA(n_components30).fit(patches_norm[train_idx]) patches_pca pca.transform(patches_norm).reshape(-1, patch_size, patch_size, 30) # 第四步NHWC - NCHWPyTorch才认 patches_pca patches_pca.transpose(0, 3, 1, 2)到了这一步patches_pca的形状就是(N, 30, 11, 11)标准化、降维、通道布局全部到位可以直接进模型了。30个主成分是我反复试下来一个比较稳的取值小于20精度有明显的下滑大于50精度提升有限训练开销却涨得很快。3. 2D_CNN网络结构设计参数不多但每个都有讲究3.1 网络整体结构我的网络结构不算新奇就是非常标准的卷积块堆叠加分类头但细节上有几个针对性设计。整体流程是输入(30, 11, 11)两组卷积块每组是两个3x3卷积堆叠再接一次2x2最大池化一个全局平均池化把特征图压成一个固定长度的向量Dropout然后接全连接层输出16类为什么用两组就可以因为11x11的输入本身不大两组卷积后的特征图尺寸约2x2再加一层也没什么信息可提了。这个网络的总参数量只有几十万级别远小于一个3D网络这是它能在小样本数据集上不炸的关键。3.2 完整代码实现import torch import torch.nn as nn import torch.nn.functional as F class HSI_CNN2D(nn.Module): def __init__(self, in_channels30, num_classes16): super().__init__() self.features nn.Sequential( # 第一组卷积 nn.Conv2d(in_channels, 64, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.Conv2d(64, 64, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 第二组卷积 nn.Conv2d(64, 128, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.Conv2d(128, 128, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), ) self.classifier nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Dropout(0.5), nn.Linear(128, num_classes), ) def forward(self, x): x self.features(x) x self.classifier(x) return x model HSI_CNN2D(in_channels30, num_classes16)这段代码可以直接跑。我特意在每个卷积后跟了BatchNorm2d它是高光谱小数据集训练的稳定器能显著抑制内部协变量偏移带来的训练抖动。3.3 几个设计决策背后的原因第一为什么卷积核全部用3x3而不是5x5或7x7两个3x3卷积串联得到的感受野等于一个5x5卷积但参数量只有后者的18/25而且中间多了一次非线性变换特征表达更强。对于11x11的小输入这个感受野尺寸刚刚好覆盖中心像素周围的纹理结构。第二为什么用AdaptiveAvgPool2d而不是直接Flatten这主要是为了兼容不同尺寸的输入。如果你之后想把patch size从11改成15或者21全连接层的输入维度会随特征图尺寸变化而崩掉。全局平均池化把任意尺寸的特征图压成固定长度这样网络结构就不用改了。第三为什么Dropout放0.5高光谱样本少过拟合是最大风险。训练集只有八千左右样本最后的全连接层又直接决定分类结果0.5的Dropout相当于强制让网络学习到更鲁棒的特征组合测试精度通常能提高2到3个点。4. 训练与评估准确率从哪来该信多少4.1 数据集划分与类别不平衡的坑Indian Pines的类别分布非常不均匀。地物类别中有的人样本超过2000个比如大豆-传统耕作有2455个样本而燕麦只有20个样本。如果随机划分训练集和测试集那几个稀有类别可能会全部落入测试集导致模型根本没学过它们。正确的做法是按类别分层抽样。每一类里面按比例抽一部分做训练剩下的做测试。常见比例是每类抽10%或20%做训练对稀有类别尤其要保证至少留下几个样本做验证。代码里需要这么做from sklearn.model_selection import train_test_split def stratified_split(labels, train_ratio0.2, seed3407): train_idx, test_idx [], [] unique_labels np.unique(labels) for cls in unique_labels: cls_idx np.where(labels cls)[0] # 稀有类别至少留2个样本进训练集防止完全没学到 if len(cls_idx) 10: n_train max(2, int(len(cls_idx) * train_ratio)) else: n_train int(len(cls_idx) * train_ratio) cls_train, cls_test train_test_split( cls_idx, train_sizen_train, random_stateseed, stratifyNone ) train_idx.extend(cls_train) test_idx.extend(cls_test) train_idx np.array(train_idx) test_idx np.array(test_idx) return train_idx, test_idx还要注意固定随机种子。我在这个项目里一直用seed3407不是玄学是为了让实验可以复现。你后面调参的时候如果每次划分都不一样你就分不清精度的变化是模型改动的功劳还是数据划分的随机波动那实验就白做了。4.2 训练超参数的选择训练参数我在多次实验后确定了一套比较稳的组合。优化器用Adam初始学习率1e-3权重衰减1e-4学习率用阶梯式衰减每个epoch乘0.95。因为Indian Pines的patch数量不大训练一个epoch很快衰减率设得平滑一点模型在后半程还能稳定打磨。import torch.optim as optim from torch.utils.data import DataLoader, TensorDataset train_dataset TensorDataset(torch.FloatTensor(patches_pca[train_idx]), torch.LongTensor(labels[train_idx] - 1)) # 标签从0开始 test_dataset TensorDataset(torch.FloatTensor(patches_pca[test_idx]), torch.LongTensor(labels[test_idx] - 1)) train_loader DataLoader(train_dataset, batch_size128, shuffleTrue, num_workers0, pin_memoryTrue) test_loader DataLoader(test_dataset, batch_size128, shuffleFalse, num_workers0, pin_memoryTrue) model HSI_CNN2D(in_channels30, num_classes16).cuda() criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3, weight_decay1e-4) scheduler optim.lr_scheduler.ExponentialLR(optimizer, gamma0.95) epochs 80 best_acc 0.0 for epoch in range(epochs): model.train() running_loss 0.0 for inputs, targets in train_loader: inputs, targets inputs.cuda(), targets.cuda() optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, targets) loss.backward() optimizer.step() running_loss loss.item() scheduler.step() # 每个epoch结束在测试集上看一下 model.eval() correct 0 total 0 with torch.no_grad(): for inputs, targets in test_loader: inputs, targets inputs.cuda(), targets.cuda() outputs model(inputs) _, predicted torch.max(outputs, 1) total targets.size(0) correct (predicted targets).sum().item() acc 100.0 * correct / total if acc best_acc: best_acc acc torch.save(model.state_dict(), best_model.pth) if (epoch 1) % 10 0: print(fEpoch [{epoch1}/{epochs}] Loss: {running_loss/len(train_loader):.4f}, Test Acc: {acc:.2f}%)batch size我用了128。Indian Pines的数据量不大显存占用不高即使只有一个中端显卡也跑得动。在RTX 3060上80个epoch大概只需要几分钟你甚至可以边训练边刷手机。4.3 评估指标OA、AA与Kappa缺一不可高光谱分类领域光报一个准确率是外行做法。标准评估要覆盖三个指标总体精度Overall Accuracy, OA所有测试样本中分类正确的比例。这个指标最直观但在类别不平衡时会被大类别主导。平均精度Average Accuracy, AA每个类别的召回率先单独算一遍再取平均。它给稀有类别同等的权重能反映模型在小类上的真实水平。Kappa系数用来衡量分类结果和随机分类相比到底好多少公式里考虑了两个类别数量不平衡带来的影响。只报OA很容易出现“总体95%但燕麦类一个都分不对”的情况这在Indian Pines上太常见了。评估代码from sklearn.metrics import accuracy_score, cohen_kappa_score, confusion_matrix def evaluate(model, X, y): model.eval() with torch.no_grad(): X_tensor torch.FloatTensor(X).cuda() outputs model(X_tensor) _, preds torch.max(outputs, 1) preds preds.cpu().numpy() y y - 1 # 转回0起始 oa accuracy_score(y, preds) cm confusion_matrix(y, preds) # 每个类别的召回率recall再取平均得到AA per_class_acc cm.diagonal() / cm.sum(axis1) aa np.mean(per_class_acc) kappa cohen_kappa_score(y, preds) return oa, aa, kappa, cm我自己在固定随机种子下用20%训练样本、30个PCA分量、11x11窗口80个epoch跑出来的结果大约在OA 96%左右、AA 93%左右、Kappa 0.95左右。不同机器、不同PyTorch版本跑出来会有一点浮动但基本在这个区间。如果你第一次跑出的结果低于90%大概率是数据预处理或训练参数里的某个环节出了偏差。4.4 分类结果可视化分类精度只是数字高光谱分类最终要落在“图”上。Indian Pines的ground truth是一张145x145的整图训练时只用有标签的像素但推理时可以对全图所有像素包括背景都做预测生成一张完整的分类图。具体做法是把整张图按滑动窗口方式切成patch包括背景像素然后逐批预测最后把预测结果填回原始的位置。可视化我用matplotlib直接用imshow就能出图。import matplotlib.pyplot as plt # 假设full_pred是(145, 145)的整图预测类别1-16背景为0 plt.figure(figsize(12, 5)) plt.subplot(1, 2, 1) plt.imshow(gt, cmaptab20) plt.title(Ground Truth) plt.axis(off) plt.subplot(1, 2, 2) plt.imshow(full_pred, cmaptab20) plt.title(Prediction) plt.axis(off) plt.savefig(hsi_result.png, dpi200, bbox_inchestight)真正有价值的观察点在边缘区域。比如林地与草地的交界处光谱曲线相似度很高人眼都很难区分2D_CNN有时候会在这里出现碎斑点。如果分类图上的噪声点太多可以考虑加大patch size或做一次简单的多数投票后处理。5. 踩坑记录与提效建议5.1 数据泄漏预处理统计量必须来源于训练集这是我最早犯过的错误也是新手最容易忽略的。有人在划分数据集之前就对全部像素做了标准化和PCA拟合然后才划分train/test。这样测试样本的均值和方差信息已经间接参与了训练数据的缩放过程测试集就不再“干净”了。测评结果虚高而且虚高的幅度在小数据集上能到3到5个百分点。我现在的固定流程是先划分索引再在训练索引上fit scaler和PCA之后把相同的transform应用到全体数据上。多写几行代码换来的是可信的实验结论。5.2 patch size不是越大越好我把patch size从11调到15、21分别做过实验。结论是15x15的准确率比11x11略高一点21x21反而下降。原因也不难理解窗口越大中心像素周围的邻域越广但Indian Pines的空间分辨率低每个像素对应的地面范围大大窗口很容易把相邻不同类别的地物包进来给模型的输入引入噪声标签。而且patch size变大后每个patch的像素数从121涨到441内存占用涨了三倍多训练速度也明显变慢。所以我最终选择了11x11作为默认配置兼顾精度和效率。如果你想调参从11和15这两个值开始试不要一上来就21。5.3 GPU相关torch.cuda.is_available()的True和False之间环境这块最常见的问题是PyTorch装了CPU版。很多时候你用Anaconda默认命令conda install pytorch装出来的就是CPU版本GPU根本没用上epoch迭代慢到怀疑人生。装完第一件事就是跑那个is_available检查。如果返回False先看你的显卡驱动NVIDIA用户在命令行执行nvidia-smi看右上角的CUDA Version。这个数字是驱动支持的最高CUDA版本你的PyTorch对应的CUDA版本必须低于等于这个数字。然后到PyTorch官网选择对应的安装命令重新安装。如果你在Windows上用WSL2NVIDIA驱动会自动映射进WSL环境不需要在WSL里重装驱动。我实际用下来WSL2里跑PyTorch的显存管理比原生Windows更听话OOM概率也小一些。5.4 想让结果更稳还能往哪个方向扩展2D_CNN跑通之后可以做三个低成本扩展第一加数据增强。旋转90度、翻转、小幅平移对地物图像来说都是安全的变换不会破坏光谱信息。用一个简单的RandomHorizontalFlip(p0.5)搭配RandomVerticalFlip(p0.5)能把有效训练样本量翻倍我实测对AA指标有1到2个百分点的提升。第二换一个更强的分类头。在全连接层之前加一个SE注意力模块让网络自动学习不同光谱通道的重要程度。这个改动很轻量但对那些光谱特征相近的类别大豆和玉米的不同耕作方式有明显帮助。第三导出ONNX做推理部署。训练完的模型可以用torch.onnx.export导出成ONNX格式之后既能用ONNX Runtime在CPU上快速推理也能转换到TensorRT做GPU加速。对于需要批量处理多景高光谱影像的实际项目来说这一步几乎是必须的。dummy_input torch.randn(1, 30, 11, 11).cuda() torch.onnx.export( model, dummy_input, hsi_cnn2d.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}} )我个人在实际操作中的体会是高光谱分类任务里先把基线做扎实比什么都重要。2D_CNN配合PCA在这里是性价比非常高的组合参数少、收敛快、结果稳特别适合作为项目的第一个里程碑。你把这个流程完整跑通一遍对高光谱图像数据和卷积网络的匹配逻辑都会有一个非常直观的理解之后再去看3D卷积或者Transformer那套方案就不会再被表面的复杂度唬住了。
网站建设高端定制企业官网