机器学习代码心得之模板和张量库

作者:陈天奇,本文首发在作者的个人微博  。
机器学习代码的很大一个特点是依赖于矩阵和向量操作,这在神经网络和矩阵分解类模型里面尤其明显。从神经网络里面的backprop到矩阵分解模型里面的更新法则都可以以向量和矩阵甚至张量的形式出现。

写机器学习程序有两类心态,一种是最大化代码运行效率的程序员,另外一种是最求实现简单少写代码的程序员。比较常见的是前者写C++,CUDA, 后者写matlab, python。我个人属于前者,希望可以对代码有最大的控制,提前分配好内存,运行中不进行内存的动态分配,如果用GPU的时候让所有的数据都呆在显卡上面不要出来等等。

当然,提到优化,最需要注意的是阿姆达法则(Amdahl'slaw),即如果你优化了占运行效率的1%的部分,再怎么优化也无济于事。很多时候数据边界检查和必要的assert尽量加上其实很多时候不会影响效率。

现在的问题是,对于极度最求效率的C++程序员,是否可以写出最求简洁明快的matlab python程序员一样的代码呢。答案是肯定的,常见的c++库Eigen和支持CUDA的mshadow就一定程度上提供了这样的功能。使用张量矩阵库的好处是可以写出简洁的代码,但是这样一来,和写pythonmatlab又有什么样的区别呢?

区别是存在的,利用c++ 的模板技术,写矩阵向量代码可以和手写C++的函数来的一样高效。这一类模板编程技术称作expressiontemplate。

具体的原因可以考虑如下问题,假设A,B,C是三个向量。 应该如何实现A = B + C这样一个操作。对于一般的矩阵库,常见的实现是重载+这个操作,分配一个新的临时空间用以储存结果,然后把B+ C的结果存到那个临时的储存空间里面去。同样的情况可以对应于其它各类操作如矩阵乘法。如果是在乎效率的C++程序员不会希望有这样的临时空间的分配,因此会直接写类似以下代码来解决问题:

for(int i = 0; i < A.len; ++i) {

A[ i ] = B[ i ]+ C[ i ];
}

问题是我们需要写很多这样的函数来对付各种的情况。而且代码也不再像A=B+C这样简洁了。利用expressiontemplate的技术,可以解决这一问题。方案如下,我们还是重载operator+这个操作,但是这个操作不进行计算,而是返回一个抽象语法表达式BinaryAddExp。然后再重载等号这个操作operator=,在等号操作的似乎我们已经拿到了目标A以及B, C。可以把实现直接写成上面的样子。这样我们可以高效地写A=B+C, 并且还是对我们的内存分配有最大的控制。样例代码见https://github.com/tqchen/mshadow/blob/master/example/exp-template/exp_lazy.cpp

当然这样的方案只能解决A=B+C,还有其它的问题如A = B+C+D或者W =gradient * eta + W 这样的更新公式。这一类更长的表达式需要抽象语法表达式的嵌套,也可以通过模板来实现。具体教程见https://github.com/tqchen/mshadow/wiki/Expression-Template。

总结下来,最重要的一点是因为这些模板展开都是在编译时执行的,使得这些复杂的组合实际上没有消耗运行时的效率。机器学习代码也可以既简洁有高效,这样的特性也只有C++模板才可以做到。对于其他语言,相对的反而不是那么自然了。

当然底层的模板库也有一些自身的缺陷,比如不可以做远距离的依赖分析,以及并行优化等。更加高层的张量库如Minerva就可以完成这些事情。但是作为一个在极度追求效率和对于资源控制的C++机器学习程序里面,基于模板的矩阵库依然是一个必不可少的组成部分。即使是最最求效率可以写公式而不是直接调用blas函数。并使得C++的机器学习代码也可以优雅高效。这相比于java也是一个重要的优势。

留下你的评论