[论文解读] Acceleration of Non-Linear Minimisation with PyTorch

Bojan Nikolic|arXiv (Cornell University)|May 18, 2018

Model Reduction and Neural Networks参考文献 9被引用 3

一句话总结

本论文展示了通过利用自动微分和GPU加速，PyTorch能够显著加速科学计算中的非线性函数最小化，仅需极少的代码修改，即可在传统NumPy实现的基础上实现高达100倍的加速。该方法利用PyTorch的自动微分和GPU卸载功能，高效计算梯度并评估复杂函数，从而在最大似然相位恢复等优化问题中实现更快的收敛速度。

ABSTRACT

I show that a software framework intended primarily for training of neural networks, PyTorch, is easily applied to a general function minimisation problem in science. The qualities of PyTorch of ease-of-use and very high efficiency are found to be applicable in this domain and lead to two orders of magnitude improvement in time-to-solution with very small software engineering effort.

研究动机与目标

• 展示原本专为神经网络训练设计的PyTorch，可有效应用于通用科学函数最小化问题。
• 通过利用PyTorch的自动微分和GPU加速，缩短非线性最小化问题的求解时间。
• 表明与传统的基于NumPy的实现相比，仅需极少的软件工程工作量，即可获得显著的性能提升。
• 在显著提升科学优化任务计算效率的同时，保持代码的可读性和正确性。

提出的方法

• 使用PyTorch的自动反向模式微分，无需手动推导或数值近似，即可计算复杂函数的梯度。
• 使用PyTorch张量实现科学函数，实现与NumPy几乎无差别的CPU和GPU执行，仅需极少代码修改。
• 通过PyTorch的CUDA支持，将张量操作卸载到NVIDIA Tesla K20c GPU上，加速数据并行计算。
• 采用标准优化算法（如BFGS），并启用jac=True标志，得益于PyTorch的自动梯度计算。
• 使用高级张量操作，以简洁易读的形式表达复杂科学模型（例如基于Zernike多项式的相位恢复）。
• 在不同网格尺寸和模型参数数量下，对PyTorch和NumPy实现进行基准测试，以衡量性能提升。

实验结果

研究问题

• RQ1 PyTorch的自动微分和GPU加速是否可有效复用于神经网络之外的一般科学函数最小化问题？
• RQ2 与传统的基于NumPy的方法相比，PyTorch在非线性最小化问题中能将求解时间缩短多少？
• RQ3 在科学优化任务中，随着模型复杂度和数据规模的增加，PyTorch的性能提升是否具有可扩展性？
• RQ4 在科学计算中，PyTorch是否能在保持高性能的同时，维持代码的可读性和可维护性？

主要发现

• 在相位恢复问题中，与仅使用CPU的NumPy实现相比，PyTorch实现了高达100倍的求解时间改进。
• 对于网格尺寸N ≥ 256的情况，GPU卸载的PyTorch实现比仅使用CPU的PyTorch版本快约10倍。
• 对于中等和大网格尺寸，仅使用CPU的PyTorch实现比NumPy实现快约一个数量级。
• 性能提升在很大程度上与模型参数数量无关，即使Zernike多项式阶数高达n=8，该现象依然成立。
• 所有实现均在算法容差范围内收敛到相同解，证实了性能提升的同时保持了正确性。
• PyTorch实现保持了与NumPy相当的代码可读性和可维护性，降低了科学模型中出现细微错误的风险。