告别2D CNN:高光谱图像分类的3D卷积原理与PyTorch实现
简介面向遥感图像分析与深度学习入门者这套资源提供基于PyTorch的3D卷积神经网络高光谱图像分类完整工程覆盖网络结构搭建、数据集划分、模型训练与推理预测等核心流程并且包含Indian Pines标准数据集与现成训练参数开箱即可复现实验。资源共11个文件主体是4个Python脚本分别承担网络定义、训练、预测与辅助工具角色另含2个.mat格式原始高光谱数据、2张结果标签图与预测图、2个Python缓存文件及1个参数文件压缩包仅23.35MB便于快速下载与二次开发。目前已有4653人学习下载社区关注度较高。通过阅读代码可以直观看到三维卷积如何同时捕获光谱与空间特征理清数据从加载、划分到训练、评估的完整链路配合预训练参数和可视化结果可快速对比超参数或网络改动带来的影响适合作为高光谱分类入门实验的起点与性能基线。1. 为什么是3D_CNN高光谱三维数据立方体带来的必然选择高光谱分类这个方向我见过太多人一上来就踩同一个坑把高光谱图像当成普通RGB图像直接套一个2D CNN模型上去折腾几天精度死活上不去。我刚开始做高光谱分类的时候也这么干过后来才发现问题不在网络深度也不在调参而是压根没用对卷积方式。要理解为什么高光谱分类绕不开3D_CNN得先看清楚手里的数据长什么样。1.1 高光谱数据的三位一体结构普通图像每个像素只有3个波段红绿蓝而高光谱图像每个像素带有几十到几百个连续波段。以最经典的Indian Pines数据集为例图像尺寸是145x145像素每个像素有200个波段整个数据就是一个145x145x200的三维立方体。其他常用数据集也类似Pavia University是610x340x103Kennedy Space Center是512x614x176。这个结构意味着每条空间位置上的一点都对应一条完整的光谱反射率曲线不同地物植被、土壤、水体、建筑的曲线形态差异往往非常明显。这个立方体直接决定了算法选型高光谱分类既要利用每个像素在光谱维的曲线形态差异又要利用空间维的邻域相关性——同一种地物的像素几乎总是聚集成块相邻像素大概率属于同一类别。也就是说分类时用到的关键信息同时分布在三个维度上任何只保留其中一到两个维度的做法本质上都在丢掉信息。1.2 2D CNN的天然短板与3D卷积的解法很多新手包括多年前的我第一反应是把高光谱数据当成一堆2D图像用2D CNN逐波段处理或者在波段维上压缩出几个统计特征再送进网络。逐波段处理的问题在于光谱维不是孤立的相邻波段之间的数值变化是连续且有物理意义的硬生生把它们拆开等于丢掉了光谱曲线形态这个高光谱数据里最关键的判别信息。而先压缩波段再分类又可能把最有区分度的波段组合给抹掉了属于一种极其粗暴的特征工程。3D CNN解决这个问题的方式非常直接——把卷积核本身升级为三维。一个形如(8, 3, 3)的3D卷积核在光谱维上同时感知连续8个波段、在空间维上感知3x3邻域滑动一个体素块就同时提取了光谱与空间的局部特征。打个比方2D CNN看高光谱像翻一本相册每次只看单独一张照片3D CNN则像在看一段连续视频单帧内容与帧与帧之间的连续变化可以同时被理解。这种空间光谱联合建模的能力就是高光谱分类任务中3D CNN效果普遍优于2D CNN的根本原因。想跑通一套完整的3D_CNN高光谱分类流程选PyTorch是当前最省心的路径生态成熟、资料多、踩坑也容易搜到答案。2. 数据准备与加载把ENVI/Mat文件变成能训练的张量写模型之前数据加载这一关就拦住了不少人。高光谱数据的存储格式五花八门不同数据集读取方式还不一样这里我把环境配置和读取细节一次说清楚。2.1 环境配置参考清单先确认环境我的建议组合是Python 3.9 PyTorch 1.13或2.x CUDA 11.7/11.8。下面是我在几台机器上验证过比较稳的搭配组件推荐版本备注Python3.8 - 3.10太新版本容易遇到个别库没跟上PyTorch1.13.0 或 2.02.x对显存优化更好CUDA11.7 / 11.82.0也支持CUDA 12.xnumpy / scipy最新即可读mat文件需要scipyspectral0.22处理ENVI格式很方便如果显存小于6G我建议先用CPU跑通整个流程再说。Indian Pines这个数据量其实CPU也能训练只是慢一些先把流程走通、把指标计算对再考虑换GPU提速。2.2 读取高光谱数据的常见格式与HDR文件解析高光谱数据常见格式有ENVI.hdr .dat/.bsq/.bil、Matlab.mat、HDF5.h5和TIFF。最容易踩坑的是ENVI格式它没有统一的读取入口需要按元数据信息手动解析。好在HDR文件已经把关键信息写得很清楚我习惯用文本方式打开它先看几个必查字段samples每行像素数宽lines图像行数高bands波段数data type数据类型1uint82int164float3212uint16interleave数据排列方式bsq按波段、bil按行、bip按像素数据量小用spectral库最方便数据量大、要按块读的时候rasterio更合适。这里给一个两种格式的读取模板import numpy as np import scipy.io as sio from spectral.io import envi # 读取 .mat 格式Indian Pines 官方数据是这种 def load_mat_data(mat_path): img sio.loadmat(mat_path) # Indian Pines 的 key 通常是 indian_pines_corrected 和 indian_pines_gt data img[indian_pines_corrected].astype(np.float32) label img[indian_pines_gt].astype(np.int64) # 0 表示背景 return data, label # 读取 ENVI 格式.hdr 和 .dat 同目录 def load_envi_data(dat_path): img envi.open(dat_path.replace(.dat, .hdr), dat_path) data np.array(img.load()).astype(np.float32) return data # 标签文件一般是单独的 .npy 或另一份 ENVI 文件注意一个隐秘的坑.mat里的数据类型可能是uint16直接丢给模型训练没问题但如果后面要做归一化、求梯度最好提前转成float32不然有些操作会因为整型溢出给出匪夷所思的结果。2.3 归一化、数据划分与Patch提取收到数据后的第一步永远是归一化。高光谱不同波段的数值范围可能差异很大不归一化会导致高数值波段主导梯度。我实践中常用两种方案一是全局min-max把整个数据立方体线性缩放到[0,1]实现最简单稳定性好二是逐波段z-score消除不同波段的量纲差异适合光谱数值分布极不均衡的数据集。我的个人经验是Indian Pines用全局min-max就够Pavia这种灰度范围差异大的逐波段归一化效果更稳。def normalize(data, modeminmax): if mode minmax: data (data - np.min(data)) / (np.max(data) - np.min(data)) elif mode zscore: mean data.mean(axis(0, 1), keepdimsTrue) std data.std(axis(0, 1), keepdimsTrue) 1e-8 data (data - mean) / std return data.astype(np.float32)样本划分是另一个关键决定。学术论文里最常见的做法是按像素随机采样比如每类随机取30%做训练、70%做测试。这种划分实现简单、类间平衡好控制但有一个严重问题——相邻像素高度相关随机采样会把同一块地物里彼此挨着的像素分别划进训练集和测试集模型实际上早就见过了测试像素的邻居指标会虚高。工程上更严谨的做法是按空间块划分或者拿同一场景不同区域的像素做测试。这个我在第5章展开讲这里先记住这个认知随机划分的结果只能作为实验内部对比不能直接当成真实泛化能力。数据准备好之后就要把原始像素转成3D_CNN需要的输入Patch。3D_CNN不只看目标像素本身还要看它周围大概多大的邻域。以Indian Pines为例我常取patch_size11或13即以目标像素为中心取一个11x11的空间邻域波段维保留全部200个波段样本维度变成(200, 11, 11)。处在图像边缘的像素没有完整邻域需要用填充策略补出来默认用reflect镜像填充比补零效果好不少因为反射填充模拟了图像边界的局部结构连续性。def extract_patches(data, labels, patch_size11, train_maskNone): 从高光谱数据中提取以每个有效像素为中心的patch。 返回的 X 形状为 (num_samples, bands, patch_size, patch_size) pad patch_size // 2 data_pad np.pad(data, ((pad, pad), (pad, pad), (0, 0)), modereflect) rows, cols, bands data.shape X, y [], [] for i in range(rows): for j in range(cols): if labels[i, j] 0: # 背景像素跳过 continue if train_mask is not None and not train_mask[i, j]: continue patch data_pad[i:i patch_size, j:j patch_size, :] X.append(patch.transpose(2, 0, 1)) # (bands, h, w) y.append(labels[i, j] - 1) # 类别从0开始 return np.stack(X), np.array(y)这一步在Indian Pines上生成的训练样本数通常是几千到上万每个样本都是(200, 11, 11)的张量内存完全放得下。3. 模型代码逐层拆解3D卷积、BN与全连接层的完整实现数据变成张量之后就到了整套代码的核心——3D_CNN模型本身。我用一个参考经典HybridSN思路设计的简洁网络来演示结构不复杂但每个组件的选择都有讲究。3.1 整体架构设计逻辑先明确输入输出输入是(1, bands, patch_size, patch_size)经过网络后输出每个类别的概率分布。整体由三组3D卷积BNReLU堆叠再接空间池化、展平、全连接和Softmax。设计时有两个关键点第一光谱维要逐步压缩。Indian Pines有200个波段第一层3D卷积核在光谱维设为8、步长为2把光谱维从200压缩到97左右第二层继续压缩第三层之后光谱维收敛到1。如果光谱维在中途不减后续特征图体积会爆炸全连接层参数也会多到难以训练。这个光谱维逐层收敛的设计是高光谱3D_CNN跟普通3D视频分类网络最大的区别需要根据输入波段数手调卷积核和步长。第二BatchNorm加在卷积和ReLU之间。网络深了之后梯度不稳定高光谱数据不同样本的光谱数值分布差异又大BN能显著加速收敛还能起到一定的正则化作用。我实测下来去掉BN的训练曲线波动幅度明显更大。3.2 核心代码完整可运行的3D_CNN类下面这个类基于PyTorch实现可以直接复制进你的训练脚本里。注释写得很细方便一行行对照理解。import torch import torch.nn as nn import torch.nn.functional as F class Hyperspectral3DCNN(nn.Module): def __init__(self, bands200, num_classes16, patch_size11, dropout0.5): super().__init__() # conv1: (光谱维8, 空间3x3)步长在光谱维上为2 self.conv1 nn.Conv3d(1, 8, kernel_size(8, 3, 3), stride(2, 1, 1), padding(0, 1, 1)) self.bn1 nn.BatchNorm3d(8) # conv2: 光谱维继续压缩 self.conv2 nn.Conv3d(8, 16, kernel_size(5, 3, 3), stride(2, 1, 1), padding(0, 1, 1)) self.bn2 nn.BatchNorm3d(16) # conv3: 空间池化之前最后一级特征提取 self.conv3 nn.Conv3d(16, 32, kernel_size(3, 3, 3), stride(1, 1, 1), padding(0, 1, 1)) self.bn3 nn.BatchNorm3d(32) self.pool nn.MaxPool3d(kernel_size(1, 2, 2)) # 动态计算展平后的特征维度 self._flatten_dim self._compute_flatten_dim(bands, patch_size) self.fc1 nn.Linear(self._flatten_dim, 128) self.fc2 nn.Linear(128, 64) self.fc3 nn.Linear(64, num_classes) self.dropout nn.Dropout(dropout) def _compute_flatten_dim(self, bands, patch_size): # 构造一个假输入前向传播得到展平后的维度避免手算错 dummy torch.zeros((1, 1, bands, patch_size, patch_size)) with torch.no_grad(): x self.pool(self.bn3(self.conv3(self.bn2(self.conv2(self.bn1(self.conv1(dummy))))))) return int(torch.flatten(x, 1).shape[1]) def forward(self, x): x F.relu(self.bn1(self.conv1(x))) x F.relu(self.bn2(self.conv2(x))) x F.relu(self.bn3(self.conv3(x))) x self.pool(x) x torch.flatten(x, 1) x F.relu(self.fc1(x)) x self.dropout(x) x F.relu(self.fc2(x)) x self.fc3(x) return x_compute_flatten_dim这个方法值得单独说一下手动算3D卷积每一层输出尺寸很容易出错让模型自己拿一个全零张量过一遍前向动态推断展平维度是最稳妥最省脑子的做法。换数据集、换patch_size都不用改代码。3.3 参数量与显存估算在Indian Pines场景下输入是(1, 200, 11, 11)第一层卷积输出约(8, 97, 11, 11)第三层输出(32, 12, 11, 11)池化后(32, 12, 5, 5)展平是9600维。全连接层第一层权重就是9600x128约123万个参数主要参数都集中在这里。一个batch 64样本做一次前向反向中间张量占用大概在1~2GB量级普通8GB显存的卡够用。如果换Pavia这种610x340的大图patch尺寸不变时训练样本更多但单个batch显存压力差异不大主要耗时在线程数。4. 训练、评估与可视化OA/AA/Kappa全套代码模型定义好之后训练和评估环节同样有许多容易忽略的细节。这一节给出直接能跑的完整训练流程和三个核心分类指标的实现。4.1 数据迭代器与训练循环构造DataLoader时shuffle只在训练集上开启验证集和测试集保持顺序即可。分类任务损失函数用交叉熵优化器用Adam初始学习率1e-3配合ReduceLROnPlateau调度。from torch.utils.data import TensorDataset, DataLoader import torch.optim as optim def get_dataloader(X, y, batch_size64, shuffleTrue): dataset TensorDataset(torch.from_numpy(X), torch.from_numpy(y)) return DataLoader(dataset, batch_sizebatch_size, shuffleshuffle) def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss, correct, total 0, 0, 0 for X_batch, y_batch in loader: X_batch, y_batch X_batch.to(device), y_batch.to(device) optimizer.zero_grad() out model(X_batch) loss criterion(out, y_batch) loss.backward() optimizer.step() total_loss loss.item() * len(y_batch) correct (out.argmax(1) y_batch).sum().item() total len(y_batch) return total_loss / total, correct / total我在超参上踩过不少次给一个宁可保守也不要激进的经验值超参数推荐值说明batch_size64或128显存小就64别硬撑learning_rate1e-3Adam RNPlateau从这里起步epochs100~200用早停防过拟合weight_decay1e-4小样本场景很有用dropout0.5全连接层前生效建议训练时记录每次epoch的loss和验证OA画一条曲线出来。如果验证曲线在前20个epoch内还往下走、之后开始横跳说明学习率需要调低如果一直不收敛优先检查数据归一化是否做了、标签是否从0开始编号。4.2 三个核心指标OA、AA、Kappa的完整实现高光谱分类论文里评价模型只看三个数总体精度OAOverall Accuracy、平均精度AAAverage Accuracy和Kappa系数。OA就是分类正确的样本占总样本的比例AA是每个类别各自精度的算术平均能反映模型在小类别上的表现Kappa衡量分类结果与随机分类相比的一致性公式相对绕一些。三个指标一起看才能全面评价模型好坏。from sklearn.metrics import confusion_matrix, cohen_kappa_score def evaluate(model, loader, device, num_classes): model.eval() all_pred, all_label [], [] with torch.no_grad(): for X_batch, y_batch in loader: X_batch X_batch.to(device) out model(X_batch) all_pred.extend(out.argmax(1).cpu().numpy()) all_label.extend(y_batch.numpy()) all_pred np.array(all_pred) all_label np.array(all_label) # OA: 总体精度 oa (all_pred all_label).mean() # AA: 各类精度的平均用混淆矩阵对角线计算 cm confusion_matrix(all_label, all_pred, labelsrange(num_classes)) per_class_acc cm.diagonal() / (cm.sum(axis1) 1e-8) aa per_class_acc.mean() # Kappa kappa cohen_kappa_score(all_label, all_pred) return oa, aa, kappa, cm, per_class_acc注意cm.sum(axis1)可能为0因为某些类别在测试集中可能一个样本都没有加一个1e-8防止除零。4.3 分类结果图可视化训练完别急着打印指标了事把整幅图预测一遍生成分类图是判断模型有没有学到东西最直观的方式。做法很简单对所有像素含背景提取patch用模型预测再把预测结果按原始坐标填回二维矩阵。背景像素单独置0否则图上会出现一块毫无意义的全预测区域。def predict_full_map(model, data, patch_size11, devicecpu): rows, cols, _ data.shape pad patch_size // 2 data_pad np.pad(data, ((pad, pad), (pad, pad), (0, 0)), modereflect) pred_map np.zeros((rows, cols), dtypenp.int64) model.eval() with torch.no_grad(): for i in range(rows): for j in range(cols): patch data_pad[i:i patch_size, j:j patch_size, :] x torch.from_numpy(patch.transpose(2, 0, 1)).float().unsqueeze(0) out model(x.to(device)) pred out.argmax(1).item() pred_map[i, j] pred 1 # 类别号从1开始对应0为背景 return pred_map有了预测图和真实标签图可以并排用matplotlib画出来也可以叠加在灰度图上观察地物边界是否符合直觉。5. 高频踩坑复盘维度报错、显存爆炸与信息泄漏最后这部分我把在实际跑这套代码过程中遇到最多、也最容易被新手忽略的四个问题整理出来。每个都是真实踩过的坑照着我说的排查思路走能省下好几个晚上的调试时间。5.1 conv3d的5维输入报错RuntimeError的常见触发点很多从2D CNN迁移过来的同学第一次跑3D_CNN会直接撞上RuntimeError: Expected 5D input for conv3d。原因很简单Conv3d要求输入形状是(N, C, D, H, W)即批量维、通道维、深度维光谱维、高、宽。而你从extract_patches得到的是(B, bands, patch, patch)只有四维少了一个通道维。修复方式是在送入网络前加一个unsqueeze(1)把(B, bands, patch, patch)变成(B, 1, bands, patch, patch)。我在前面的代码里已经默认了输入有这个维度但很多从网上拷代码的人会把这一步漏掉。遇到这个报错先检查输入张量维度不要急着怀疑网络结构。5.2 显存爆炸patch_size、batch_size和波段数的三角博弈高光谱分类特别容易爆显存原因是每个patch都保留了全部波段200个通道的中间特征图比普通3通道图像膨胀了几十倍。同一个batch下patch_size从11涨到15显存占用不是线性涨而是接近立方增长。我试过8G显存的卡跑patch_size15、batch_size64前向一次直接OOM。解决的顺序依次是先降batch_size到32或16再考虑降patch_size最后才动网络宽度。另外一个很实用的方案是使用PyTorch 2.x的自动混合精度在损失计算前调用torch.cuda.amp.autocast()显存能省下差不多三分之一精度损失基本可以忽略。5.3 指标虚高随机划分带来的信息泄漏问题这个坑最隐蔽也最影响结果可信度。按像素随机划分训练集和测试集时目标像素周围邻域里的像素很可能也被划进了训练集而3D_CNN输入的patch就包含这些邻域。模型在训练时已经见过测试样本的邻居空间结构测试精度自然偏高。在Indian Pines上这种泄漏带来的虚高有时能达到几个百分点。如果你想评估模型的真实泛化能力建议至少做一次按空间块划分的对比实验把整幅图切成互不重叠的区块一部分区块做训练另一部分做测试这样patch之间没有重叠指标才是实打实的。5.4 验证集精度震荡与过拟合训练曲线一直上下乱跳最常见的两个原因一是学习率太大二是batch太小导致BatchNorm统计量不稳定。我的排查习惯是先把学习率降到1e-4跑20个epoch如果曲线平滑了说明是学习率问题如果还是震荡把batch_size提到128再看。小样本过拟合是另一个高频问题Indian Pines有些类别只有几十个样本模型很容易把训练集背下来。我验证过比较有效的组合是全连接层前加Dropout(0.5) 训练时对patch做随机翻转和随机裁剪 weight_decay开到1e-4。这三个手段叠在一起能把验证集和训练集精度之间的差距拉小很多。最后分享一个我自己实验里的习惯训练前固定随机种子包括numpy、random和PyTorch的都要设不然每次跑出来的结果都不一样调参时根本不知道是改超参数的功劳还是运气。另外新手跑这类高光谱项目我强烈建议先拿小patch_size比如9把整个流程跑通确认指标计算正确、分类图可视化正常再换大patch追求精度。这个顺序看起来慢实际是最快的路径。本文还有配套的精品资源点击获取