基于ONNX与测算一体的端侧轻量化洪水预报

ONNX-Based Lightweight On-Device Flood Forecasting with Integrated Measurement and Computation

摘要: 洪涝灾害频发且破坏性强，对洪水预报的时效性、经济性与场景适应性提出更高要求。针对传统洪水预报模型计算量大、云端部署存在传输延迟高、运维成本高、偏远地区适配性差等问题，本研究提出一种基于开放神经网络交换(ONNX)的端侧轻量化LSTM测算一体的洪水预报方法。以山东省内6个AiFlow视觉测流监测站点的流量数据为研究基础，通过Yeo Johnson转换将偏态分布的流量数据转换为近似正态分布，结合Z-score标准差标准化构建时序训练样本；采用双层堆叠LSTM模型捕捉水文数据的时序依赖关系，基于ONNX对模型开展计算图优化、算子融合与量化处理，实现模型的端侧轻量化部署与河道流量实时滚动预报。实验结果表明，该方法在6个监测站点均表现出可靠的预报精度，1、3步滚动预报纳什系数(NSE)均高于0.91，平均绝对百分比误差(MAPE)均小于10%，满足短期洪水预报需求；经ONNX轻量化处理后，较原始PyTorch格式模型，模型文件体积减少了67%，可执行文件体积减少了69%，端侧推理冷启动耗时降低超99%，6个站点的平均推理耗时降低7%~31%，模型预报NSE无变化；模型加载内存增量从111.97 MB降低至0.54 MB，完全适配端侧设备的资源约束。本方法实现了“数据采集、本地推理、预报输出”的全流程闭环，显著降低了洪水预报的运维成本和预报延迟，为中小河流、偏远地区的测算一体化洪水预报提供了低成本、轻量化的技术方案。

Abstract: Frequent and highly destructive flood disasters have imposed increasingly stringent requirements on the timeliness, cost-effectiveness, and scenario adaptability of flood forecasting. To address the key limitations of traditional flood forecasting models (namely excessive computational overhead) and the inherent drawbacks of cloud-based deployment (including high transmission latency, elevated operation and maintenance (O&M) costs, and poor adaptability to remote areas), this study proposes a lightweight on-device LSTM (Long Short-Term Memory) flood forecasting method with integrated measurement and computation based on the Open Neural Network Exchange (ONNX). Taking the discharge data from 6 AiFlow visual flow measurement monitoring stations in Shandong Province as the research basis, the skewed discharge data are transformed into an approximately normal distribution via the Yeo-Johnson transformation, and time-series training samples are constructed in conjunction with Z-score standardization. A two-layer stacked LSTM model is built to capture the temporal dependencies of hydrological series, and ONNX-based computation graph optimization, operator fusion, and quantization processing are implemented on the trained model, enabling lightweight on-device deployment of the model and real-time rolling forecasting of river channel discharge. The experimental results demonstrate that the proposed method achieves reliable forecasting performance across all 6 monitoring stations. For 1-step and 3-step ahead rolling forecasting, the Nash-Sutcliffe Efficiency (NSE) values are all above 0.91, and the Mean Absolute Percentage Error (MAPE) values are all below 10%, fully satisfying the operational requirements of short-term flood forecasting. After ONNX-enabled lightweight processing, the model code is reduced by 67%, the executable file size is reduced by 69%. The cold start latency of on-device inference is cut by over 99%, and the average inference time across the 6 stations is reduced by 7% to 31%, with no degradation in forecasting NSE. The memory increment required for model loading is reduced from 111.97 MB to 0.54 MB, which fully complies with the resource constraints of on-device edge terminals. The proposed method establishes a full-process closed loop covering “data acquisition, local inference, and forecasting output”, thus significantly lowers the O&M costs and forecasting latency of flood forecasting, and provides a low-cost, lightweight technical solution for integrated measurement-computation flood forecasting in medium and small rivers and remote areas.