获取符号矩阵微分的计算图/表达式

Get computational graph / expression of symbolic matrix differentiation

我想为神经网络编写一些自定义 CUDA 内核以加快计算速度,但如果有可以自动执行的程序包,我不想花时间手动微分张量表达式。

是否有python包可以显示符号矩阵微分的表达式?

我知道 sympy 可以像这样处理非矩阵表达式:

def func(x):
    return 1 / x

arg_symbols = sp.symbols(inspect.getfullargspec(func).args)
sym_func = func(*arg_symbols)
s = ''
for arg in arg_symbols:
    s += '{}\n'.format(arg, sp.Lambda(arg_symbols, sym_func.diff(arg)))
# this is what I need:
print(s)
>>> Lambda(x, -1/x**2)

我知道autograd包可以计算矩阵表达式的导数

After the function is evaluated, autograd has a list of all operations that were performed and which nodes they depended on. This is the computational graph of the function evaluation. To compute the derivative, we simply apply the rules of differentiation to each node in the graph.

但是有没有办法从它或一些类似的包中得到这个微分计算图?

您引用的包之间存在一些严重差异。差异是您无法直接从自动微分库中获取(AFAIK)计算图的原因,但您可以从基于符号的库中获取它。

简而言之:

  • 数值微分:numpy就够了
  • 符号微分:sympy
  • 自动微分:autograd为一例

区分方法分为三种:

  • 数值微分:计算 Delta(f(x)) / Delta(x),其中 Delta(x) 是一个小差异,代表 x 的变化,同时保持在 f 的范围内。这不是你需要的。你不需要这个包。
  • symbolic differentiation:基于表示函数符号应用的图的构造(我在Ruby here中有一篇关于符号引擎实现的文章).在这种情况下,微分是通过链推导规则的递归应用来执行的:

    f(g(x))' = f'(g(x)) * g'(x)
    

    当这个规则应用到整个符号图时,结果是一个更新的带有导数的符号图。优点在于导数是精确的,但是对于非常复杂的图,最终的导数图可能是不可处理的(超过内存限制或堆栈限制太深的递归)。 在python中sympy实现了这种推导。另一方面,如果您有导数图,您可以对其执行操作,例如简化或替换。

    from sympy import *
    import numpy as np
    
    x = symbol('x')
    f = 1 / x
    df = diff(f, x)
    print(df)
    # -1/x**2
    
    ldf = lambdify((x), df)
    
    # Now ldf is a lambda
    x_ary = np.array([
      [[1, 2, 3], [1, 2, 3]], 
      [[1, 2, 3], [1, 2, 3]]
    ])
    y_ary = ldf(x_ary)
    
    print(xn.shape)
    # (2, 2, 3)
    print(y_ary)
    # array([[[-1.        , -0.25      , -0.11111111],
    #         [-1.        , -0.25      , -0.11111111]],
    #        [[-1.        , -0.25      , -0.11111111],
    #         [-1.        , -0.25      , -0.11111111]]])
    

    如您所见,它与 numpy 一起使用,但它涵盖了一些基本示例而不是所有内容,实际上 sympy.matrixsympy.symbol 应该用于特定图形(例如:我认为它不能直接处理diff(x.T A x, x) = x.T A + A x))。

    也可以 export the graph as C code,但它的功能有限,对于您的应用程序,您一定要修改结果:

    from scipy.utilities.codegen import codegen
    
    [(cf, cc), (hf, hc)] = codegen(("df", df), "C", "df")
    
    print(hc, cc)
    

    打印出来:

/*****************************************************
 *      Code generated with sympy 1.1.1              *
 *  See http://www.sympy.org/ for more information.  *
 *      This file is part of 'project'               *
 *****************************************************/

#ifndef PROJECT__DIFF__H
#define PROJECT__DIFF__H

double df(double x);

#endif

/*****************************************************
 *      Code generated with sympy 1.1.1              *
 *  See http://www.sympy.org/ for more information.  *
 *      This file is part of 'project'               *
 *****************************************************/
#include "diff.h"
#include <math.h>

double df(double x) {   
  double df_result;
  df_result = -1/pow(x, 2);
  return df_result;   
}
  • 自动微分就是通过autograd完成的。在这种情况下,两全其美。从一方面来说,没有必要显式评估图形,另一方面,您不能在派生函数上执行进一步的操作,同时保持导数的精确性。这是(通常)通过使用附加字段(类似于 float[2] 或更多)扩充 float 定义来完成的,其中附加字段包含导数。例如,在自动微分环境中,sin 函数可能会重载:

    def sin(x):
         return [sin(x[0]), x[1] * cos(x[0])]
    

    但是你可以这样理解,没有可用的计算图,而是直接得到精确导数沿x的值(所有函数必须重载)。我有一个更完整的示例(在 C 语言中,仅使用宏)here请注意,Tensorflow 内部使用自动微分而不是符号微分,但建议用户直接提供 "explicit version" 来处理数值不稳定性!。自动微分通常不处理数值不稳定性。