ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

基于CNN的水下声源定位实战:从PyTorch训练到Jetson部署

基于CNN的水下声源定位实战:从PyTorch训练到Jetson部署 简介海洋水下声源定位是水声工程与深度学习结合的热点任务这套基于PyTorch的CNN实战资源面向声源方位角与距离预测适合高校本科生毕设、研究生课题起步、课程设计及科研原型验证尤其适配人工智能、水声工程、电子信息、自动化等方向学生。压缩包共59个文件、约7.24MB以22个Python脚本为主体覆盖数据生成、模型训练与推理演示另含11张PNG网络结构示意图、2个Jupyter Notebook、MATLAB脚本及详细README便于交互阅读和二次开发。项目已在Windows 10/11和macOS实测通过支持CPU与GPU一键运行目前已有33人学习下载。资源提供完整原始声学数据、数据生成脚本、训练验证代码和已收敛的预训练权重无需从零搭建即可运行demo查看预测效果也可替换数据或调整网络层结构开展实验。配套README.md对环境配置、数据加载、模型训练、推理演示及常见问题均有说明可作为海洋感知场景中深度学习落地的实用参考。 水下声源定位这个方向听起来挺高端但本质上就是让机器回答一个问题某个声音是从哪个位置发出来的。放在海洋工程里它的应用非常实际——无人潜航器避碰、水下目标监测、海洋哺乳动物科考、海底管线巡检都要依赖声源定位。我这次把整套流程从数据到部署完整跑了一遍把PyTorch版CNN模型、完整数据集、预训练权重和部署指南全部整理成了一整套可直接复用的实战项目。这里把我的踩坑经历、设计思路和关键代码逐段拆开讲清楚希望能给做水声信号处理、深度学习结合落地的朋友省下几周时间。1. 项目整体设计与方案选型1.1 为什么用CNN做水声定位传统的水声定位思路主流方案是波束形成Beamforming、互相关时延估计TDOA这类信号处理方法。它们本身很成熟但落地时有两个绕不开的问题第一对信噪比要求高海洋环境噪声复杂浅海混响、舰船辐射噪声、生物叫声叠在一起传统算法性能掉得很快第二多途效应导致信号到达水听器阵列时已经不是简单的几何延迟叠加而要经过复杂的信道畸变传统模型很难把这种非线性关系建模好。CNN的方案思路完全不同。我把它理解成一种“数据驱动的地图匹配”阵列接收信号的时频特征本身携带了声源位置的信息网络通过大量样本学习“特征—位置”之间的映射关系。它不需要你显式建模声传播信道只要训练数据分布覆盖了实际场景卷积层就能自动提取出跟位置强相关的空间线索。实测下来在中等信噪比5~15 dB条件下CNN定位精度明显优于传统互相关方法特别是多途干扰严重时优势更明显。1.2 技术选型与全局架构整套系统我用PyTorch 2.x作为深度学习框架主要考虑三点动态图机制方便快速验证网络结构TorchScript和ONNX导出链路成熟训练完可以直接部署到嵌入式设备HuggingFace、TensorBoard等生态工具集成度高。信号处理部分用了NumPy SciPy做仿真数据生成和预处理比直接上Librosa更轻因为水下声信号处理更关注包络和相位信息Librosa很多面向音频的功能用不上。系统总体分五个模块数据仿真与预处理按设定声源位置、信道参数生成水听器阵列接收信号提取时频特征PyTorch数据集封装自定义Dataset类实时做数据增强避免一次性全部加载到内存CNN模型多分支输入结构提取空间特征和时间特征后融合全连接层输出坐标训练与验证AdamW优化器 余弦退火学习率调度用RMSE和CEP指标评估部署训练后的模型导出成ONNX再转TensorRT引擎部署到Jetson Orin NX板卡2. 数据集构建不准备仿真后面全白搭2.1 仿真数据生成与关键参数水下声源定位数据集不像图像数据集在网上一抓一大把真实海试数据获取成本极高而且很难获得准确的声源位置真值Ground Truth。所以这个项目的数据集是仿真的但仿真不是乱造参数完全参考了浅海环境的标准模型。我用的参数如下阵型4元十字水听器阵列阵元间距0.75 m对应中心频率1 kHz的半波长声速1500 m/s典型海水声速采样率48 kHz信号类型CW脉冲和线性调频信号LFM两种频率范围0.5~2 kHz声源距离50~500 m声源方位角0~360度信噪比-5~20 dB 随机分布信道模型加入多途延迟、瑞利衰落、海洋环境噪声近似高斯分布 低频风成噪声仿真的核心代码思路是先把声源位置转换成到每个阵元的真实距离差得到时延差再根据多途模型生成几路叠加回波最后加上不同信噪比的噪声。总共生成10万条样本按8:1:1划分训练/验证/测试集。这里有个容易踩的坑——如果只在高信噪比下仿真模型训练完到实际场景基本瘫痪所以信噪比范围一定要拉大。2.2 预训练权重怎么用项目里提供了预训练权重这背后有个很实用的迁移学习思路。我最初在纯仿真数据上训练的模型拿到湖试实测数据上一测精度掉了不少。原因很简单仿真数据和真实数据之间存在域差异仿真信道的简化模型没法完全覆盖真实环境的声传播特性。解决方案是两阶段训练第一阶段在10万条仿真数据上训练100轮让模型学到通用的声场特征表示第二阶段用少量实测数据几十条带真值的样本做微调只训练最后两层全连接层学习率降到1e-4。这样既保留了仿真数据学到的大规模先验又能快速适配实测数据的分布。项目里提供的预训练权重就是第一阶段训练出来的你现在有少量实测数据的话可以直接加载它做微调不用从零开始训练。3. 模型架构与PyTorch实现细节3.1 网络结构设计思路定位模型我采用了多分支CNN结构可以理解成这样一个流程每个阵元接收到的信号先各自提取时频图然后通过共享权重的卷积骨干网络提取特征最后把所有特征融合起来回归坐标。这个设计的出发点很直接每个阵元都是一个独立的“传感器”它们各自对声音到达方向和时延有独立的感知能力而声源位置的关键信息恰恰藏在阵元之间的差异里所以特征融合阶段不能简单做加和要做特征拼接后再让全连接层自己学习阵元间的关系。模型的具体结构是这样的输入4路时频图每路尺寸128×256频率×时间每路分支3层Conv2D BatchNorm ReLU MaxPool特征融合4路特征拼接后过一层SE注意力模块回归头两层全连接 GELU激活输出2维x, y坐标卷积核统一用3×3通道数依次为32、64、128。一开始我试过更大的256通道效果提升很少但参数量和训练时间翻倍所以128就足够了。SE注意力模块在这个任务里特别好用它相当于让网络学会“哪路阵元的特征更可信就多给它一点权重”在多途干扰下某些阵元信号可能被严重污染注意力机制能自动压制这些不可靠信息。3.2 关键实现代码模型核心代码大致是这样的import torch import torch.nn as nn import torch.nn.functional as F class SeBlock(nn.Module): def __init__(self, channels, reduction16): super().__init__() self.squeeze nn.AdaptiveAvgPool2d(1) self.excitation nn.Sequential( nn.Linear(channels, channels // reduction), nn.ReLU(inplaceTrue), nn.Linear(channels // reduction, channels), nn.Sigmoid() ) def forward(self, x): b, c, _, _ x.size() w self.squeeze(x).view(b, c) w self.excitation(w).view(b, c, 1, 1) return x * w class HydrophoneBranch(nn.Module): def __init__(self, in_channels1, out_channels128): super().__init__() self.features nn.Sequential( nn.Conv2d(in_channels, 32, 3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, out_channels, 3, padding1), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue), nn.AdaptiveAvgPool2d(4), ) def forward(self, x): return self.features(x) class UnderwaterLocator(nn.Module): def __init__(self, num_hydrophones4, n_hidden256): super().__init__() self.num_hydrophones num_hydrophones self.branch HydrophoneBranch() self.se SeBlock(num_hydrophones * 128) self.head nn.Sequential( nn.Flatten(), nn.Linear(num_hydrophones * 128 * 16, n_hidden), nn.GELU(), nn.Dropout(0.3), nn.Linear(n_hidden, 2), ) def forward(self, x): # x: (B, N, C, H, W) N为阵元数 b, n, c, h, w x.shape x x.view(b * n, c, h, w) feats self.branch(x) feats feats.view(b, n, -1, 4, 4) feats feats.view(b, n * 128, 4, 4) feats self.se(feats) out self.head(feats) return out有个细节值得注意4路阵元的卷积层是共享权重的也就是说只有一个HydrophoneBranch实例被复用了四次。这么做的好处是大大减少参数量而且本质上让每个阵元用相同的特征提取逻辑避免了某一路阵元因为训练样本不均衡被特殊对待。损失函数这块我分别试过MSE和Smooth-L1Huber Loss。MSE收敛快但对离群点太敏感有时个别严重污染的样本会把训练带偏。Smooth-L1在高误差区域梯度是常数训练更稳定。最终选的是Smooth-L1并且在坐标输出前加了Sigmoid约束到[0,1]再用归一化系数映射到实际坐标范围这样避免网络在最开始训练时就预测出离谱的大坐标。4. 训练配置、模型评估与部署全流程4.1 训练参数与评估体系训练配置直接抄作业优化器AdamW初始学习率1e-3权重衰减1e-4学习率调度余弦退火CosineAnnealingLR最小学习率1e-5Batch Size64训练轮次150轮数据增强随机时移±50个采样点、随机加减信噪比±3 dB、随机丢弃一路阵元信号模拟阵元故障硬件单张RTX 3090训练耗时约2.5小时有人会问定位这种回归任务为什么要做随机丢弃阵元这种“缺胳膊少腿”的数据增强我的体会是实际布放阵列时水听器出现故障、线路接触不良是常有的事如果模型没有见过“丢阵元”的输入运行时一路信号断了预测结果会完全放飞。加了这种增强后模型对单路阵元失效的容忍度明显提高这是实测中非常实用的处理方式。评估指标不能只看一个。平均定位误差RMSE代表整体精度但水声定位领域大家更关心CEP圆概率误差含义是50%的定位结果落在多大半径的圆内。我测试集上的结果如下信噪比区间RMSE米CEP50米15~20 dB3.22.15~15 dB7.85.4-5~5 dB18.512.7可以看到信噪比高的时候定位精度能到米级这个精度够无人潜航器做近程避碰和目标跟踪吗其实很勉强通常还需要后续的卡尔曼滤波平滑轨迹但作为单帧定位已经算是很能打的水平。4.2 从PyTorch模型到嵌入式部署模型训练完不部署就是“实验室模型”我的落地目标是Jetson Orin NX这类嵌入式平台。部署链路是PyTorch → ONNX → TensorRT FP16。PyTorch转ONNX的代码不复杂但有几个坑要提醒import torch from model import UnderwaterLocator model UnderwaterLocator() model.load_state_dict(torch.load(best_model.pt, map_locationcpu)) model.eval() dummy_input torch.randn(1, 4, 1, 128, 256) torch.onnx.export( model, dummy_input, locator.onnx, input_names[hydrophone_inputs], output_names[xy_coords], dynamic_axes{hydrophone_inputs: {0: batch}, xy_coords: {0: batch}}, opset_version17, )第一坑dynamic_axes一定要加否则后续TensorRT推理时batch size被锁死。第二坑BN层和Dropout层在导出前必须切换成eval()模式否则权重是乱的。第三坑如果模型里用了torch.where这种动态控制流算子ONNX导出大概率失败我的方案是避开这些算子改成masked_fill或者where的乘加实现。ONNX导出后转TensorRT在Jetson上我用的是TensorRT 8.6版本FP16精度批量大小为1。转换命令trtexec --onnxlocator.onnx --saveEnginelocator_fp16.engine --fp16实测下来RTX 3090上PyTorch推理一帧约3.2毫秒Jetson Orin NX上FP16 TensorRT推理一帧约8.5毫秒完全满足实时性需求。从FP32换FP16精度几乎没有损失定位误差只增加了0.3米左右这在工程上完全可接受。部署端我直接用Python调用TensorRT的Python API CUDA Stream一条流水线从信号采集卡读数据、预处理、推理、输出到显控界面整体延迟控制在20毫秒以内。5. 常见问题与避坑经验5.1 训练阶段的高频问题问题一损失函数已经收敛但验证集定位误差突然在某个信噪比下特别差。这个现象我排查了很久最后发现是训练数据里低信噪比样本占比太少模型对低信噪比场景欠拟合。解决方法是按信噪比分层采样保证每个批次里都有各个信噪比区间的样本。顺带说一句训练数据的信噪比分布最好和实际使用场景对齐否则无论你怎么调网络结构都是“巧妇难为无米之炊”。问题二模型在仿真数据上精度不错但一拿到实测数据就崩溃。这基本就是域差异问题我前面说的两阶段微调就是干这个用的。这里有一个进阶技巧微调时可以把卷积层的前几层冻结requires_gradFalse只微调高层特征和回归头。因为底层卷积学到的是时频图的局部纹理特征这在不同域之间是通用的而高层特征和位置映射关系才是需要适配的部分。问题三训练过程中偶尔出现loss突然跳变到NaN。我遇到的原因是输入时频图中个别样本的全部数值为0阵列信号采集时有一段静默经过归一化后出现除零。解决办法是在数据预处理时对全零序列做保护加一个极小值eps1e-8。5.2 部署阶段的高频问题问题一ONNX导出成功但TensorRT构建引擎时报错“op not supported”。这个八成是模型里用了不常见算子。GELU激活函数在较老版本TensorRT就不支持我的解决方式是直接用ReLU替代GELU精度几乎没有损失或者在转ONNX时将GELU展开成数学表达式Sigmoid近似形式。问题二FP16推理结果在某些信噪比下偏差大。FP16的精度动态范围约在1e-3到6e4之间如果输入特征数值本身分布范围很宽FP16量化会导致小数值信息丢失。我的方案是在预处理阶段对特征做全局归一化让数据落在[0,1]区间内。实测归一化后FP16和FP32的精度差距明显缩小。问题三阵列通道数量不一致时的推理问题。项目输出4元阵列的模型但有些应用场景只有2元或者3元阵列。我的处理方法是训练阶段做随机通道遮罩mask把4路输入随机丢弃若干路推理阶段如果是3元阵列就补一路全零输入给模型。实测3元阵列定位性能大概比4元差20%~30%但仍然可用这也算是模型对传感器故障的鲁棒性扩展。写在最后我跑完整个项目最深的体会是水下声源定位的难点不在网络结构有多花哨而在数据和部署的工程化。仿真数据设计直接影响模型上实测的表现部署链路的每个环节都可能因为算子兼容性问题卡住。目前这套项目已经把数据集、预训练权重、PyTorch源码和部署脚本全部打包好了拿到手就可以从数据可视化、模型训练到板卡部署完整跑通。后续如果要做多声源定位或移动声源跟踪只需要把输出头改成多目标回归或者在模型后面接一层卡尔曼滤波这个架构做底子是够用了。本文还有配套的精品资源点击获取
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进