利用PyTorch JIT快速训练LSTM

Posted 2021-05-02 图灵TOPIA

图灵TOPIA

编辑：刘静

图灵联邦编辑部出品

JIT（Just-In-Time）是一组编译工具，用于弥合 PyTorch 研究与生产之间的差距。它允许创建可以在不依赖 Python 解释器的情况下运行的模型，并可以更积极地进行优化。

训练 RNN 和 LSTM 是非常困难的，因为计算能力受到内存和带宽等的约束。简而言之，每个 LSTM 单元需要 4 个线性层（MLP 层），以便每个顺序时间步运行一次。线性层需要大量的内存带宽才能执行计算；由于系统没有足够的内存带宽将数据馈送到计算单元，实际上它们无法使用许多计算单元。

近期，Thomas Viehmann发文介绍了他的相关研究成果，可以提高LSTM的训练速度。

Thomas Viehmann的身份是培训数学家、专业顾问精算师、夜间机器学习业余爱好者、MathInf的创始人。（MathInf是一个机器学习和精算咨询机构）

以下是相关介绍：

实施快速递归神经网络(RNN)是一项具有挑战性的任务。这不仅是因为训练现有架构很麻烦（有时是对CuDNN的LSTM优化实现很麻烦）。更糟糕的是，它还限制了对新架构的实验。

fast.ai的Jeremy Howard写道：

例如，我们一直在对不同类型的RNN架构和层规范化进行大量研究。在这两种情况下，即使使用PyTorch出色的新JIT编译器，我们也无法获得与纯CUDA C实现相同的性能水平。

他的评论暗示了解决这个问题的传统方法：编写（更好的）CUDA kernel。

我已经为Batch Norm和CTC Loss做了这项工作（现在PyTorch上的速度与我GPU上的CuDNN速度相当，cuDNN是NVIDIA深度神经网络软件开发包中的其中一种加速库）。

但这需要付出很多努力，并且它并不能真正用于快速实验（特别是对RNN而言）。所以你最终还是得编写了两个kernel，因为你无法使用PyTorch的自动区分。

Jeremy的结论是用新语言编写一个新框架。他不太喜欢Python，我也不喜欢Python 的一些缺点，即使我仍然认为Python是一种很棒的语言。因此承诺就是：新的Swift for TensorFlow最终将带有一个新的中间表示MLIR，这使得各种优化变得容易。

但PyTorch确实有一个很棒的JIT IR，它是我们已经看到的优化的原因。因此，与其发明一种新语言，不如努力让JIT编译器变得更加好用。

通过融合相关技术细节，我们提出一系列优化方法（补丁），使用这些补丁，PyTorch LSTM基准测试的jit-premul LSTM落后于CuDNN所需时间的1.33倍。当前进和后退时，我们比CuDNN慢约25％。 这就是用Python / PyTorch实现的LSTM单元。

我们向后加速了大约2.25倍。当我学习计算机科学时，我最喜欢的教授之一是ArnoldSchönhage，他因大整数的快速乘法而闻名。在他的《Fast Algorithms》一书中，他提出的黄金法则其中之一便是第2条规则：永远不要浪费2倍！。

所以我们似乎在这里做了有用的工作。还有第8条规则：快速算法的发展很慢！这同样适用于上述优化。

技术细节请参考：

https://lernapparat.de/fast-lstm-pytorch/#fn:branch

https://github.com/t-vi/pytorch/tree/sumtosize_fusedconcat