一种提高深度算子网络模型推理求解和训练速度的方法

本发明属于人工智能及高性能计算技术领域，特别涉及一种提高深度算子网络模型推理求解和训练速度的方法。本发明针对PI‑DeepONet模型进行多层次的优化，提高其推理求解速度和训练速度。本发明定制SGEMM算子，在PI‑DeepONet常用的模型尺寸范围内，SGEMM算子相较于CUDA计算平台的基础线性代数库中最好的单精度矩阵乘法可取得最高1.5倍的加速比。针对模型中使用的modified MLP的特殊结构设计融合算子，模型的端到端推理速度获得2倍以上的加速比，图形处理器计算时间最高可被缩短5.7倍。本发明使用数据并行对模型进行分布式训练，并使用通信掩盖的策略，可扩展性保持在83%‑85%。