This paper presents TorchNWP, a compilation library tool for the efficient coupling of artificial intelligence components and traditional numerical models. It aims to address the issues of poor cross-language compatibility, insufficient coupling flexibility, and low data transfer efficiency between operational numerical models developed in Fortran and Python-based deep learning frameworks. Based on LibTorch, it optimizes and designs a unified application-layer calling interface, converts deep learning models under the PyTorch framework into a static binary format, and provides C/C++ interfaces. Then, using hybrid Fortran/C/C++ programming, it enables the deployment of deep learning models within numerical models. Integrating TorchNWP into a numerical model only requires compiling it into a callable link library and linking it during the compilation and linking phase to generate the executable. On this basis, tangent linear and adjoint model based on neural networks are implemented at the C/C++ level, which can shield the internal structure of neural network models and simplify the construction process of four-dimensional variational data assimilation systems. Meanwhile, it supports deployment on heterogeneous platforms, is compatible with mainstream neural network models, and enables mapping of different parallel granularities and efficient parallel execution. Using this tool requires minimal code modifications to the original numerical model, thus reducing coupling costs. It can be efficiently integrated into numerical weather prediction models such as CMA-GFS and MCV, and has been applied to the coupling of deep learning-based physical parameterization schemes (e.g., radiation, non-orographic gravity wave drag) and the development of their tangent linear and adjoint models, significantly improving the accuracy and efficiency of numerical weather prediction.
翻译:本文提出TorchNWP,一种用于高效耦合人工智能组件与传统数值模型的编译库工具。该工具旨在解决Fortran开发的业务化数值模型与基于Python的深度学习框架之间跨语言兼容性差、耦合灵活性不足及数据传输效率低的问题。基于LibTorch,优化设计了统一的应用层调用接口,将PyTorch框架下的深度学习模型转换为静态二进制格式并提供C/C++接口,进而通过混合Fortran/C/C++编程实现深度学习模型在数值模型内部的部署。将TorchNWP集成至数值模型仅需将其编译为可调用链接库,并在编译链接阶段进行链接即可生成可执行文件。在此基础上,在C/C++层面实现了基于神经网络的切线性与伴随模型,能够屏蔽神经网络模型的内部结构,简化四维变分资料同化系统的构建过程。同时支持在异构平台上部署,兼容主流神经网络模型,可实现不同并行粒度的映射与高效并行执行。使用该工具仅需对原数值模型进行极少量的代码修改,从而降低耦合成本。该工具可高效集成至CMA-GFS、MCV等数值天气预报模式,并已应用于基于深度学习的物理参数化方案(如辐射、非地形重力波拖曳)的耦合及其切线性与伴随模型的研发,显著提升了数值天气预报的精度与效率。