随着气象业务现代化和精细化的发展，高时空分辨率的格点天气要素数据逐步建立并得到应用，高精度、格点化的气象资料不仅是天气预报和气候研究等方面的工作基础，而且在精细工农业、生态系统研究等方面也有着广泛的应用前景。对于空间精度来说，栅格点分辨率需要从当前的0.0083×0.0083经纬度（平均面积约为1平方公里）提升到0.001×0.001经纬度；对于时间精度来说，由预测1天到可以预测14天，整体时空精度的提升对算力的需求需要1000x的提升。在此情况下，只有引入高性能计算芯片才能满足在规定时间内做出及时的预测。下面我们向大家详细介绍一下通过在Graphcore IPU上加速Kriging算法，整体的计算时间在单个M2000上相对于双核CPU（原有算法只能利用单核计算，CPU型号为Intel Xeon E5-2680v3，主频率为2.5GHZ）提升了60.5倍。项目源码在文末参考资料中。

ET₀介绍

ET₀（Reference evapotranspiration）的完整概念是参考蒸发蒸腾量，其主要流程如图1所示，通过对温度、湿度、气压以及风速在内的天气要素进行插值分析和预测计算得到更高密度的气象要素点，然后对温度进行校准，使用彭曼公式算出区域内所有网格点的ET₀，最后提取出所需区域内的ET₀值并存储在数据库中。ET₀对气象数据挖掘、数值天气预报水平的评估以及气候变化等相关研究具有重要价值，也广泛应用于极端天气预测、干旱预测、智慧灌溉等领域。例如在智慧灌溉方面，ET₀是水文循环的重要组成部分，对作物需水量和水资源管理产生重大影响，准确的ET₀估算有助于正确确定水资源预算和分配，从而提高灌溉系统的用水效率。由于气象要素主要是国家气象站依靠卫星和雷达进行观测采集，气象站站点的分布范围和密度有限，站点覆盖密度大的区域要素数据精度较高，密度低的区域精度较低，因而需要把离散的监测站观测数据通过合适的空间插值方法转变成具有较高时空分辨率的空间网格点数据。由图1我们可以看到Kriging插值算法是ET₀计算的主体部分，也是我们接下来优化的主要目标。

Kriging算法详解

Kriging算法是依据协方差函数对随机过程或随机场进行空间建模和预测（插值）的回归算法。在特定的随机过程中，例如固有平稳过程中，Kriging算法能够给出最优线性无偏估计，因此在地统计学也被称为空间最优无偏估计器，被广泛应用于地理科学、环境科学、大气科学等领域。空间网格点越精细，插值预测结果越准确，但同时也会平方级增加插值算法的计算量，如果不提升算力，算法在实际应用过程中会有一些局限。

这里我们讨论的是基于线性模型的二维普通Kriging插值算法，参考了PyKrige的Python开源代码。

1. 总体算法流程

Kriging算法总体计算公式为（1）：

其中，z(S_i)是第i个位置的测量值，在本项目中就是已给定的经纬度上的温度（湿度、风速等，共2167个点），n是点临域内实测点的个数，将其展开成矩阵之后形式如下所示（2）：

其中是点（i, j）处z的半方差，其计算方法为（3）：

2. 半变异函数

从步骤1中可知计算预测点的z值需要事先知道向量和矩阵，对于已知的气象数据，由于其z值是事先给定的，其矩阵直接通过式（2）就可计算出来。对于需要预测的气象数据，暂时只知道预测点的经纬度坐标，因此需要通过该坐标得到未知点的矩阵。

因此可事先计算出已知点间的欧式距离矩阵，然后通过最小二乘法拟合这些点的半方差矩阵得到slope和nugget的值，对于需要预测的点，计算出其距离矩阵后通过式（4）即可算出预测点的半方差矩阵。

3. λ向量求解

通过步骤2可以得到已知点的半方差矩阵A以及未知点的半方差矩阵B，将其带入式（2）中即可算出 λ向量，

然后把向量带入式（1）即可算出最后的预测值z。

4. 算法流程

算法主要包含的部分如图3所示，输入数据x、y、z分别代表经度、纬度和天气要素，使用仿射变换对经纬度进行一次纠正，计算纠正后各点间的距离及对应天气要素的半方差矩阵，并利用最小二乘法得到距离与半方差矩阵间的映射关系，接着使用kd树找到距离预测点最近的n个输入点【此处的n与式（1）中的n一致，这里设置为12】，分别计算预测点与预测点间以及预测点与上述n个输入点间的距离矩阵，再利用映射参数将其转换为天气要素的半方差矩阵a与b【对应于式（5）中的A和B】，通过求逆运算计算出λ，最后带入式（1）得到最后的插值预测结果。

5. 计算耗时分析

通过上述分析可知算法主要包含矩阵间的距离计算、最小二乘拟合计算、kd树聚类计算以及矩阵点乘和求逆运算，当输入数据点为2167，预测点数据为7346*4534，每个点由附近最近的12个点负责预测时，各部分的计算规模和在CPU上的耗时如表1所示。可看出CPU端的耗时主要体现在kd树和插值预测过程中，欧式距离和拟合过程几乎不耗时，因此我们针对kd树和插值预测进行相关优化（图3中红色虚线框使用CPU实现，蓝色虚线框使用IPU实现）。

类型	输入	输出	耗时（秒）
欧式距离	2167		0.12
线性拟合		2	0.21
kd树	7346*4534, 2167, 12	7346*4534*12	34
插值预测（求逆及点乘）	7346*4534*12, 7346*4534*12	7346*4534*12	300

算法

Kriging算法在IPU上的优化

我们在IPU上实现并加速了Kriging算法，也取得了非常出色的结果。在M2000的IPU服务器上，当天气要素站点数量为2167，预测点数量为7346*4534时，IPU的执行效率是ArcGIS的15.7倍。若同时计算日平均气压、平均风速、平均湿度、平均温度、最高温度和最低温度这6个天气要素来生成完整的天气预测图ET₀，经优化后IPU的执行效率是ArcGIS的27.3倍。使用4块IPU同时计算14天内的ET₀，IPU执行效率是ArcGIS的60.5倍。主要优化过程如下：

1. kd树优化

在PyKrige的开源算法中，我们了解到kd树的实现使用的是单线程，而IPU搭载的服务器CPU核心数在50个以上。综合测试后，我们使用了8个线程对kd树进行求解，充分发挥多核优势将kd树的耗时降至8秒左右，较大程度上提升了kd树的计算效率。

2. 插值预测优化

在PyKrige的开源算法中，对于向量的求解是逐行进行的，共需执行7346*4534次，且每行的向量线性求解速度较慢，在原版Python代码中计算完所有数据耗时约2000多秒，使用Cython能极大加快计算效率，但统计下来耗时依然约300秒左右，与ArcGIS内置的速度相当。我们使用TensorFlow将求解和点乘过程移植至IPU上，并充分利用IPU的片上内存，一次执行batch_size为20000的矩阵求逆和点乘，借助于IPU强大的矩阵运算能力和大批次处理，算完所有的数据耗时只需21秒，极大提高了运行效率。

另外，我们还使用了自定义算子实现了矩阵的线性求解，跳过更加耗时的求逆运算，利用比求逆复杂度更低的矩阵LU分解方法直接对向量λ进行求解，并手动给数据分配IPU上的tile计算单元降低数据与数据间的交互和同步，进一步提升运算效率，同样的数据量大小计算耗时约为11秒。经测试，我们在IPU上对Kriging插值算法优化后的运行时间合计约为21秒，ArcGIS和PyKrige开源算法的运行时间为340秒，加速比约为16倍。

3. 预处理优化

原始的天气要素数据需要用IDL软件做一些预处理，我们对IDL的脚本做了新的实现，将预处理结果直接输入到Kriging插值算法中，整个算法是一个端到端的过程，无需存取中间计算结果；

4. 算法流程优化

由于每一天各个要素的气象点位置是一致的，输入点的欧式距离计算和kd树求解不会因为天气要素的改变而改变，所以一天内的6个天气要素在CPU端的结果完全一致，只需计算一次，从而优化了距离和kd树求解时间。

5. 多IPU并行计算

我们采用4块IPU同时计算4天的插值结果，大大提高了运行效率。

总结

IPU计算一天的ET₀总体性能如表2所示，加速比性能如表3所示，最终输出的结果比对结果如图4所示（左图为IPU实现的结果，未取出mask，右图为ArcGIS计算的结果），可看出IPU复现的结果与ArcGIS生成的天气要素图分布几乎一致。

原始数据预处理	网格点生成	kd树求解	最小二乘拟合	插值预测（6要素串行）	ET0计算	合计
0.1	0.1	8	4.9	65	7.3	85.4

参考链接

ET₀已经在Graphcore GitHub上开源，您可以在此获取代码和教程。

关于安捷中科CEO黄思源完整讲述ET₀与Graphcore合作请观看完整视频。