MindSpore AI科学计算系列（5）：AI框架加速海洋数值模拟

本篇将介绍MindSpore团队与清华大学黄小猛老师共同开发的MindSpore版GOMO（Generalized Operator Modelling of the Ocean）海洋模型。

背景介绍

( X; g% T) K$ g: \# D. Z

GOMO模型是一个区域海洋模式，最早由清华大学黄小猛老师基于OpenArray框架开发。海洋模式是指通过一组物理方程来描述海洋的气候变化，不仅可以很好地表征海面温度和高度分布，还能够实时预测台风、海啸等现象。海洋模式自1967年诞生以来发展迅速，至今已经有40多个海洋模式版本。其中比较有代表性的就有全球海洋模式（modular ocean model, MOM）：POP , 以及区域海洋模式（regional ocean model system, ROMS）：POM等。GOMO模型中的基本方程和实现算法便来自于POM模型。

传统数值计算方法如有限差分法将海洋特定区域离散成网格点，每个网格点內可以计算出流体速度、水体温度、盐度等物理量。在过去的几十年中，研究人员已经开发了许多模型来提高仿真结果（Bonan和Doney；Collins等人；Taylor等人）。这些模型变得越来越复杂，代码量已经从几千行增加到数万甚至数百万行。在软件工程方面，代码量的增加使模型更难以开发和维护。基于此，黄小猛老师设计了OpenArray并行计算库，将复杂的并行计算与海洋模式研究解耦。

基于OpenArray的GOMO海洋模式

5 p- _& O* A9 z

OpenArray将常用的差分算子进行了抽象：

研究人员可以快速方便的将离散的PDE方程转换为相应的运算符表达形式。同时这些算子在底层可以实现自动并行，用户在上层实现的串行代码和并行代码一致，从而使得研究人员免于实现复杂的并行编程。下图是OpenArray中实现海洋海表高度求解的过程：

0 b3 I6 V8 Q) _# | j5 k

OpenArray实现代码即方程（海表高度方程）

- N8 i5 y1 }1 M$ R6 W

尽管如此，OpenArray+GOMO目前仍然存在一些问题：

/ d- Q: `! Q. m9 ?5 U7 D

第一个问题是计算效率。变量在计算过程中一旦被加载到处理器的寄存器或高速缓存中，则应在替换变量之前尽可能多地使用它，频繁的变量加载和不可避免的缓存丢失，会带来极高的内存消耗，导致计算性能下降。GOMO当前的效率和可扩展性已接近sbPOM（POM模行的变种）的一系列优化方法，例如内存池，图形计算，JIT编译和向量化，这些方法都是用于降低对内存带宽的需求，提升性能。但是，OpenArray目前尚未完全解决内存带宽限制问题。

第二个问题是当前的OpenArray版本不支持自定义运算符。当用户尝试另一种高阶对流方程或任何其他数值方程时，OpenArray提供的12个基本运算符可能不能完全实现求解过程。

6 d9 e- \' \6 l

第三个问题是对硬件平台的支持，目前OpenArray只支持传统CPU集群与神威太湖之光，但是不能使用GPU和Ascend等平台。

7 `6 d, M+ [0 e) q$ |$ e; T8 y

第四个问题是无法采用自动微分的功能进行模型参数优化以及数据同化。

8 j& R8 @8 k8 X7 K* i' \5 z

MindSpore加速GOMO求解

2 V* E) v* h. R& [6 N- D! g" z

针对OpenArray+GOMO中目前存在的一些问题，可以使用深度学习框架MindSpore结合GPU对GOMO进行进一步的加速求解。

! t# ` l* m6 _$ o! K

算子抽象

借鉴OpenArray的思想，我们在MindSpore中进行类似的算子抽象。以DYF算子为例，使用Pad、Slice算子组合实现DYF运算，如下图所示。首先对输入的A(x, y)的y轴向后扩充一维，用Pad填充为0；再使用Slice算子将y轴第一维移除，这样得到的A’(x, y)中的每个元素与原始输入中的A(i, j+1)一一对应。最后使用A’(x, y)减去A(x, y)就得到了DYF算子的运算结果。

' U0 F$ l( j& u6 }4 r/ Q: k

' a5 @8 x$ C1 L: Q& i, S! k

MindSpore实现差分算子抽象

& B, Y) s8 b; J* l0 m( c, M( _

目前在已经在MindSpore中实现了12个用于平均和差分运算的算子，满足绝大多数偏微分方程的求解。同时如果用户需要进行额外的差分运算，可以参照上述的算子抽象方式，实现灵活的算子定义。

图算融合

7 {" C/ P8 r4 Q8 Z' M

图算融合是MindSpore独具特色的性能优化技术。通过自动分析和优化现有计算图逻辑，并结合目标硬件能力，对计算图进行计算化简和替代、算子拆分和融合、算子特例化编译等优化，实现对网络性能的整体优化。相比传统优化技术，图算融合具有多算子跨边界联合优化、与算子编译跨层协同、基于Polyhedral的算子即时编译等独特优势。另外，图算融合只需要用户打开对应配置后，整个优化过程即可自动完成，不需要网络开发人员进行其它额外感知，使得用户可以聚焦网络算法实现。

如下图所示，是海洋模式中求解正压模态的海表高度方程的图算融合过程。首先，MindSpore会将用户的实现代码转换为对应的计算图，用户的输入和计算过程对应计算图中的每个节点。然后未使能图算融合的原始计算图输入到AKG（Auto Kernel Generator， 自动算子生成）模块，AKG会对输入的计算图进行扫描，自动生成对应的融合算子，减少中间变量的产生，增加指令集发射长度，提高计算效率。当前MindSpore可以自动的对add、sub、mul、div这些ElementWise算子进行融合，无需用户进行额外的操作。而我们的最终目标是将自定义的差分算子与基础算子进行更大范围的融合，最终融合成一个完整的算子。这样对于每一个求解方程来说，都是一个融合算子，没有额外的中间计算结果产生，能够极大的提升性能。

, c! i6 h+ J' G# K7 s, V, @, C

4 J! {( t, p; y% M, { L) o0 a

MindSpore图算融合

/ z1 S/ Y: [ a- P) k7 M- s

我们测试了基于MindSpore实现的海洋区域模型GOMO单机版，在开启图算前后的性能对比，如下图。从测试结果可以看出，使能图算融合之后GOMO模型的单步迭代时间提升约1倍，并且对于不同的分辨率均有效果。

8 J* v* z0 r* D3 n3 ?; A7 k0 y

& K4 I& s' M' X5 \& X; u! ?: V

图算融合前后性能对比

案例介绍

( l: o5 x2 O J0 j% f; o

下面将简单介绍MindSpore GOMO模型使用。实践前先确保已经正确安装MindSpore。如果没有，可以通过MindSpore安装页面安装；其次安装netCDF4

' T- J% [( l% S1 [' C! N8 r2 q

pip install netCDF4

. Z" E/ { l7 ]3 B

1. 准备数据

+ s [ _# W# v8 u# N

本案例使用的是netCDF格式的Seamount文件，贝克曼和海德沃格尔提出的Seamount问题是区域海洋模型广泛使用的理想试验案例（Beckmann and Haidvogel, 1993）。

2. 加载数据

加载Seamount数据文件，从文件脚本中读取变量的初始化值，Seamount文件中的数据类型是双精度Float64，需要将其转成Float32进入MindSpore计算。加载处理数据的脚本在源码的src/read_var.py脚本中。

import numpy as np ! G; s4 Z. F# p3 c7 h: u. K import netCDF4 as nc # variable name list params_name = [z, zz, dz, dzz, dx, dy, cor, h, fsm, dum, dvm, art, aru, arv, rfe, rfw, : I' U! A# F# U7 A' d* B rfn, rfs, east_e, north_e, east_c, north_c, east_u, north_u, east_v, north_v, tb, . V* T! j( Y) N& N4 u/ \8 c sb, tclim, sclim, rot, vfluxf, wusurf, wvsurf, e_atmos, ub, vb, uab, vab, elb, etb, dt, uabw, uabe, vabs, vabn, els, eln, ele, elw, ssurf, tsurf, tbe, sbe, sbw, tbw, tbn, tbs, sbn, sbs, wtsurf, swrad] def load_var(file_obj, name): ( l/ w2 V: M# b; [( S$ X3 s """load variable from nc data file""" 8 k; x7 a7 x3 A& z1 J) p data = file_obj.variables[name] 7 s5 z3 A" B+ m5 W- w data = data[:] data = np.float32(np.transpose(data, (2, 1, 0))) return data " t+ ^* h0 |8 g, a) s, z+ N def read_nc(file_path): """ put the load variable into the dict """ 5 T, U5 M; c# T0 R variable = {} file_obj = nc.Dataset(file_path) 3 @9 _ n3 q" ]2 \ for name in params_name: # o( g% F0 \2 U variable[name] = load_var(file_obj, name) 1 i8 g* I1 x# B6 S- s/ U return variable

3. 定义GOMO网络

0 U& h; b- z6 p

GOMO模型基于动量、能量和质量守恒定律，推导微分方程组和边界条件，确定需要求解的7个方程组，详细的公式推导参考论文。图1是GOMO的整体执行流程图。首先，从Seamount数据中加载数据，用于模型中变量的初始化。加载初始值和模型参数后，计算分为内模态循环和外模态循环两个部分。在外模态循环中，主要计算二维海表面高度el和二维平均风速ua、va。在内模态循环中，循环次数iend是训练的总时间步数（由用户输入设定），内模态循环的计算三维数组占主导地位，依次计算湍流动能q2和产生湍流动能的湍流长度q2l、温度t和盐度s、x和y方向的风速u和v。计算完成之后，保存所需的变量结果，结束训练。

GOMO模型流程图

初始化变量

... from src.GOMO import GOMO_init ... p4 s; g$ A8 t/ R; L9 @( p if __name__ == "__main__": % ?- ^; u! r0 e3 o7 I- Q1 i+ y, d ... # define grid and init variable update net_init = GOMO_init(im, jm, kb, stencil_width) % E% |5 E' X8 Y% w ...

. O& j& ]& k7 {- @( T5 ]

定义GOMO模型

$ o+ ^3 e1 S- s5 [ }. _" T

def construct(self, etf, ua, uab, va, vab, el, elb, d, u, v, w, kq, km, kh, q2, q2l, tb, t, sb, s, rho, wubot, wvbot, ub, vb, egb, etb, dt, dhb, utb, vtb, vfluxb, et): """construct""" x_d, y_d, z_d = self.x_d, self.y_d, self.z_d ; S$ O) X& g9 u$ k q2b, q2lb = self.q2b, self.q2lb dx, dy = self.dx, self.dy ! L7 a ~; w$ u # surface forcing # v* W" d, `0 e7 I* Y' C w = w * (1 - self.z_h) + self.z_h * self.vfluxf # lateral_viscosity ( b- L& P7 O3 k advx, advy, drhox, drhoy, aam = self.lateral_viscosity(dx, dy, u, v, dt, self.aam, ub, vb, x_d, y_d, z_d, rho, self.rmean) 9 B6 L$ y C: W! l+ G" V # mode_interaction adx2d, ady2d, drx2d, dry2d, aam2d, advua, advva, egf, utf, vtf = self.mode_interaction(advx, advy, drhox, drhoy, aam, x_d, y_d, d, uab, vab, ua, va, el) # ===========external model=========== vamax = 0 elf = 0 / g+ D- c8 O) i7 O for iext in range(1, 31): * b/ X1 U5 R! e$ I% E3 q2 I% a # external_el 9 v6 i6 |# C; c3 S4 v3 Q$ n* G elf = self.external_el(x_d, y_d, d, ua, va, elb) # external_ua advua, uaf = self.external_ua(iext, x_d, y_d, elf, d, ua, va, uab, vab, el, elb, advua, aam2d, adx2d, drx2d, wubot) # external_va " S2 W! j6 E% c% U7 n u: l, T7 { advva, vaf = self.external_va(iext, x_d, y_d, elf, d, ua, va, uab, vab, el, elb, advva, aam2d, ady2d, dry2d, wvbot) # external_update etf, uab, ua, vab, va, elb, el, d, egf, utf, vtf, vamax = self.external_update(iext, etf, ua, uab, va, vab, el, elb, elf, uaf, vaf, egf, utf, vtf, d) # m6 i2 A c9 P+ ^9 b1 r8 j$ W # ===========internal model=========== & o6 l* s+ I. p9 u. Y/ F# G if self.global_step != 0: . D$ H4 ^$ R( X9 A # adjust_uv 5 X0 T4 f( _" H/ o4 Y9 n* o4 T( p u, v = self.adjust_uv(u, v, utb, vtb, utf, vtf, dt) # internal_w w = self.internal_w(x_d, y_d, dt, u, v, etf, etb, vfluxb) + }+ b/ S- F0 Y$ ` # internal_q 4 b% `% W, A- n' D, c t: V @ dhf, a, c, gg, ee, kq, km, kh, q2b_, q2, q2lb_, q2l = self.internal_q(x_d, y_d, z_d, etf, aam, q2b, q2lb, q2, q2l, kq, km, kh, u, v, w, dt, dhb, rho, wubot, wvbot, t, s) ' N* ~: j) E1 D! S0 J8 U/ r q2b = ops.Assign()(self.q2b, q2b_) q2lb = ops.Assign()(self.q2lb, q2lb_) # internal_t_t ) @, G' B/ W, Z/ ]# ]* b a, c, ee, gg, tb, t = self.internal_t_(t, tb, self.wtsurf, self.tsurf, self.swrad, self.tclim, self.tbe, self.tbw, self.tbn, self.tbs, x_d, y_d, z_d, dt, u, aam, self.h, self.dum, v, self.dvm, w, dhf, etf, a, kh, self.dzz, c, self.dzz1, ee, gg, dx, self.dz, dy, self.fsm, dhb) # internal_t_s a, c, ee, gg, sb, s = self.internal_t_(s, sb, self.wssurf, self.ssurf, self.swrad0, self.sclim, self.sbe, self.sbw, self.sbn, self.sbs, x_d, y_d, z_d, dt, u, aam, self.h, self.dum, v, self.dvm, w, dhf, etf, a, kh, self.dzz, c, self.dzz1, ee, gg, dx, self.dz, dy, self.fsm, dhb) 0 E7 A `6 r: w # dense rho = self.dens(s, t, self.zz, self.h, self.fsm) ! ?5 ?2 ]$ u1 y1 `) @4 i # internal_u uf, a, c, gg, ee, wubot = self.internal_u(x_d, z_d, dhf, u, v, w, ub, vb, egf, egb, ee, gg, self.cbc, km, advx, drhox, dt, dhb) # internal_v 7 K! u8 x: ^/ T, { vf, a, c, gg, ee, wvbot = self.internal_v(y_d, z_d, dhf, u, v, w, ub, vb, egf, egb, ee, gg, self.cbc, km, advy, drhoy, dt, dhb) # adjust_ufvf u, v, ub, vb = self.adjust_ufvf(u, v, uf, vf, ub, vb) # internal_update egb, etb, dt, dhb, utb, vtb, vfluxb, et = self.internal_update(egf, etb, utf, vtf, etf, et) ; N9 B3 E1 L8 y4 C steps = ops.AssignAdd()(self.global_step, 1) ; |/ G/ D. V* S: f return elf, etf, ua, uab, va, vab, el, elb, d, u, v, w, kq, km, kh, q2, q2l, tb, t, sb, s, rho, wubot, wvbot, \ ub, vb, egb, etb, dt, dhb, utb, vtb, vfluxb, et, steps, vamax, q2b, q2lb

; u' b( g* z" I

在__main__函数中调用定义好的GOMO模型：

. J, }0 ^! y! g: [' L W6 H

... from src.GOMO import GOMO ... if __name__ == "__main__": ... # define GOMO model Model = GOMO(im=im, jm=jm, kb=kb, stencil_width=stencil_width, variable=variable, x_d=x_d, y_d=y_d, z_d=z_d, E0 |7 E8 t! R q2b=q2b, q2lb=q2lb, aam=aam, cbc=cbc, rmean=rmean) ...

4. 训练网络

I! J6 m# N4 b$ L( x+ Z

运行脚本

" c; {, w- F5 s; h( x8 \2 c* M0 s |

训练脚本定义完成之后，调用scripts目录下的shell脚本，启动训练进程。 使用以下命令运行脚本：

sh run_distribute_train.sh <im> <jm> <kb> <step> <DATASET_PATH>

脚本需要传入变量im、jm、kb、step、DATASET_PATH，其中：

3 B- s. s2 w6 j* {# _1 A

· im,jm,kb：模拟的海洋区域分辨率，与使用的数据相关；

· step：训练的时间步数（与图1中的iend对应）；

- ?; y T# Y4 t0 K' v

· DATASET_PATH：训练数据路径。

! s g( g" c: v7 a# {! H0 R

训练完后，训练过程中变量的变化值保存在train/outputs目录下，每隔5个时间步保存一次数据，主要保存了4个变量值，分别是东向的风速、北向的风速（单位是m/s），位温度（单位是K），海表面高度（单位是m）。

; u9 d4 v6 R, K8 G& ?

└─outputs . o5 z8 F. [! k; j; N* P! F ├─u_5.npy ├─v_5.npy ; L2 V! s3 a5 h* x9 G ├─t_5.npy ├─et_5.npy 0 O/ C. }+ [' x1 { ├─u_10.npy ; ?4 ^, B5 z/ \4 a9 B! S ├─v_10.npy ; }6 [ C8 B5 E, R/ K- J ├─t_10.npy ├─et_10.npy

其中，*.npy：指保存的变量。文件名称具体含义：变量名称_step数.npy。

+ N, F. K( n1 E5 c' z6 g% s u

展望

MindSpore版本的GOMO模型，通过Python前端完成了关键差分算子的抽象，提升了易用性；同时结合图算融合功能+GPU硬件对GOMO模型进行了加速。不仅如此，用户还可以借助MindSpore的自动微分功能实现模型参数调优以及数据同化。在此，也欢迎广大的科学计算爱好者和研究者加入我们，共同拓展和维护MindSpore版本GOMO模型。

$ T) h V; t( o) V) t2 M0 ~! j* ~

参考文献

1. Huang X, Huang X, Wang D, et al. OpenArray v1. 0: a simple operator library for the decoupling of ocean modeling and parallel computing[J]. Geoscientific Model Development, 2019, 12(11).

+ |. F: ?- D: I7 a

2. Blumberg A F, Mellor G L. A description of a three‐dimensional coastal ocean circulation model[J]. Three‐dimensional coastal ocean models, 1987, 4: 1-16.

; R' \& V/ i q) D m* N3 P

3. Beckmann A, Haidvogel D B. Numerical simulation of flow around a tall isolated seamount. Part I: Problem formulation and model accuracy[J]. Journal of Physical Oceanography, 1993, 23(8): 1736-1753.

2 F" Y6 g8 F. g " ^4 m5 _) _/ C1 u ' U" S% e- c( L9 b4 X$ |' A 4 p% w9 i% x% \$ a