TensorFlow可微分编程实践3---向量微分和Jacobian矩阵

在这篇博文中，我们将利用TensorFlow Eager Execution API来实现一个完整多层感知器（MLP）模型。在具体实现多层感知器模型之前，我们首先来看，怎样用TensorFlow Eager Execution API来求向量与矩阵运算的导数。
我们知道在多层感知器模型中，最基本的运算是由第l−1$l-1$层输出信号求出第l$l$层神经元的输入信号，公式如下所示：

zl=Wlal−1+bl(3.2.001)$$\begin{array}{}\text{(3.2.001)}& {\mathit{z}}^l=W^l{\mathit{a}}^{l-1}+{\mathit{b}}^l\end{array}$$

为了下面讨论方便，我们假设第l−1$l-1$层有3个神经元，第l$l$层有2个神经元，式(3.2.001)中的各个值定义如下。
第l−1$l-1$层输出信号：

al−1=⎡⎣⎢1.02.03.0⎤⎦⎥(3.2.002)$$\begin{array}{}\text{(3.2.002)}& {\mathit{a}}^{l-1}=\left[\begin{array}{c}1.0\\ 2.0\\ 3.0\end{array}\right]\end{array}$$

第l−1$l-1$层到第l$l$层连接权值矩阵：

Wl=[4.07.05.08.06.09.0](3.2.003)$$\begin{array}{}\text{(3.2.003)}& W^l=\left[\begin{array}{ccc}4.0& 5.0& 6.0\\ 7.0& 8.0& 9.0\end{array}\right]\end{array}$$

第l$l$层偏置值：

bl=[1001.01002.0](3.2.004)$$\begin{array}{}\text{(3.2.004)}& {\mathit{b}}^l=\left[\begin{array}{c}1001.0\\ 1002.0\end{array}\right]\end{array}$$

为了进行学习，我们需要求出以下导数：∂zl∂al−1$\frac{\mathrm{\partial }{\mathit{z}}^l}{\mathrm{\partial }{\mathit{a}}^{l-1}}$、∂zl∂bl$\frac{\mathrm{\partial }{\mathit{z}}^l}{\mathrm{\partial }{\mathit{b}}^l}$、∂zl∂Wl$\frac{\mathrm{\partial }{\mathit{z}}^l}{\mathrm{\partial }{W}^l}$，我们分别来进行讨论。
我们首先来看第一项，根据Jacobian矩阵定义得：

∂zl∂al−1=⎡⎣⎢⎢⎢⎢⎢⎢⎢⎢⎢⎢∂zl1∂al−11∂zl2∂al−11...∂zlNl∂al−11∂zl1∂al−12∂zl2∂al−12...∂zlNl∂al−12............∂zl1∂al−1Nl−1∂zl2∂al−1Nl−1...∂zlNl∂al−1Nl−1⎤⎦⎥⎥⎥⎥⎥⎥⎥⎥⎥⎥=⎡⎣⎢⎢⎢⎢⎢Wl1,1Wl2,1...WlNl,1Wl1,2Wl2,2...WlNl,2............Wl1,Nl−1Wl2,Nl−1...WlNl,Nl−1⎤⎦⎥⎥⎥⎥⎥=Wl(3.2.005)$$\begin{array}{}\text{(3.2.005)}& \frac{\mathrm{\partial }{\mathit{z}}^l}{\mathrm{\partial }{\mathit{a}}^{l-1}}=\left[\begin{array}{cccc}\frac{\mathrm{\partial }{z}_1^l}{\mathrm{\partial }{a}_1^{l-1}}& \frac{\mathrm{\partial }{z}_1^l}{\mathrm{\partial }{a}_2^{l-1}}& ...& \frac{\mathrm{\partial }{z}_1^l}{\mathrm{\partial }{a}_{N_{l-1}}^{l-1}}\\ \frac{\mathrm{\partial }{z}_2^l}{\mathrm{\partial }{a}_1^{l-1}}& \frac{\mathrm{\partial }{z}_2^l}{\mathrm{\partial }{a}_2^{l-1}}& ...& \frac{\mathrm{\partial }{z}_2^l}{\mathrm{\partial }{a}_{N_{l-1}}^{l-1}}\\ ...& ...& ...& ...\\ \frac{\mathrm{\partial }{z}_{N_l}^l}{\mathrm{\partial }{a}_1^{l-1}}& \frac{\mathrm{\partial }{z}_{N_l}^l}{\mathrm{\partial }{a}_2^{l-1}}& ...& \frac{\mathrm{\partial }{z}_{N_l}^l}{\mathrm{\partial }{a}_{N_{l-1}}^{l-1}}\end{array}\right]=\left[\begin{array}{cccc}{W}_{1,1}^l& W_{1,2}^l& ...& W_{1,N_{l-1}}^l\\ W_{2,1}^l& W_{2,2}^l& ...& W_{2,N_{l-1}}^l\\ ...& ...& ...& ...\\ W_{N_l,1}^l& W_{N_l,2}^l& ...& W_{N_l,N_{l-1}}^l\end{array}\right]=W^l\end{array}$$

我们接下来再来求对第l$l$层偏置求微分：

∂zlbl=⎡⎣⎢⎢⎢⎢⎢⎢⎢⎢⎢⎢∂zl1∂bl1∂zl2∂bl1...∂zlNl∂bl1∂zl1∂bl2∂zl2∂bl2...∂zlNl∂bl2............∂zl1∂blNl∂zl2∂blNl...∂zlNl∂blNl⎤⎦⎥⎥⎥⎥⎥⎥⎥⎥⎥⎥=⎡⎣⎢⎢⎢10...001...0............0０...1⎤⎦⎥⎥⎥(3.2.006)$$\begin{array}{}\text{(3.2.006)}& \frac{\mathrm{\partial }{\mathit{z}}^l}{{\mathit{b}}^l}=\left[\begin{array}{cccc}\frac{\mathrm{\partial }{z}_1^l}{\mathrm{\partial }{b}_1^l}& \frac{\mathrm{\partial }{z}_1^l}{\mathrm{\partial }{b}_2^l}& ...& \frac{\mathrm{\partial }{z}_1^l}{\mathrm{\partial }{b}_{N_l}^l}\\ \frac{\mathrm{\partial }{z}_2^l}{\mathrm{\partial }{b}_1^l}& \frac{\mathrm{\partial }{z}_2^l}{\mathrm{\partial }{b}_2^l}& ...& \frac{\mathrm{\partial }{z}_2^l}{\mathrm{\partial }{b}_{N_l}^l}\\ ...& ...& ...& ...\\ \frac{\mathrm{\partial }{z}_{N_l}^l}{\mathrm{\partial }{b}_1^l}& \frac{\mathrm{\partial }{z}_{N_l}^l}{\mathrm{\partial }{b}_2^l}& ...& \frac{\mathrm{\partial }{z}_{N_l}^l}{\mathrm{\partial }{b}_{N_l}^l}\end{array}\right]=\left[\begin{array}{cccc}1& 0& ...& 0\\ 0& 1& ...& ０\\ ...& ...& ...& ...\\ 0& 0& ...& 1\end{array}\right]\end{array}$$

下面是一个向量对矩阵求偏导，而我们对这个操作没有定义，所以我们需要以一种变通的方式来进行，我们将Wl$W^l$视为由w(i)∈RNl−1${\mathit{w}}^{(i)}\in R^{N_{l-1}}$的行向量组成：

Wl=⎡⎣⎢⎢⎢⎢(w(1))T(w(2))T...(w(Nl))T⎤⎦⎥⎥⎥⎥=wl∈RNl(3.2.007)$$\begin{array}{}\text{(3.2.007)}& W^l=\left[\begin{array}{c}({\mathit{w}}^{(1)}{)}^T\\ ({\mathit{w}}^{(2)}{)}^T\\ ...\\ ({\mathit{w}}^{(N_l)}{)}^T\end{array}\right]={\mathit{w}}^l\in R^{N_l}\end{array}$$

其实w(i)${\mathit{w}}^{(i)}$是指向第l$l$层第i$i$个神经元所有连接权值组成的向量。
有了式(3.2.007)的定义，我们就可以将Wl$W^l$视为向量，这样根据Jacobian矩阵定义：

∂zl∂Wl=∂zl∂wl=⎡⎣⎢⎢⎢⎢⎢⎢⎢⎢zl1w(1)zl2w(1)...zlNlw(1)zl1w(2)zl2w(2)...zlNlw(2)............zl1w(Nl)zl2w(Nl)...zlNlw(Nl)⎤⎦⎥⎥⎥⎥⎥⎥⎥⎥(3.2.008)$$\begin{array}{}\text{(3.2.008)}& \frac{\mathrm{\partial }{\mathit{z}}^l}{\mathrm{\partial }{W}^l}=\frac{\mathrm{\partial }{\mathit{z}}^l}{\mathrm{\partial }{\mathit{w}}^l}=\left[\begin{array}{cccc}\frac{z_1^l}{{\mathit{w}}^{(1)}}& \frac{z_1^l}{{\mathit{w}}^{(2)}}& ...& \frac{z_1^l}{{\mathit{w}}^{(N_l)}}\\ \frac{z_2^l}{{\mathit{w}}^{(1)}}& \frac{z_2^l}{{\mathit{w}}^{(2)}}& ...& \frac{z_2^l}{{\mathit{w}}^{(N_l)}}\\ ...& ...& ...& ...\\ \frac{z_{N_l}^l}{{\mathit{w}}^{(1)}}& \frac{z_{N_l}^l}{{\mathit{w}}^{(2)}}& ...& \frac{z_{N_l}^l}{{\mathit{w}}^{(N_l)}}\end{array}\right]\end{array}$$

与前面不同的是，式(3.2.008)的矩阵中的每个元素都是一个标量对向量的求偏导，根据我们上篇博文介绍，标量对向量求偏导，结果为一个行向量，我们以zliw(j)$\frac{z_i^l}{{\mathit{w}}^{(j)}}$为例进行讨论。
如果i≠j$i\ne j$时，w(j)${\mathit{w}}^{(j)}$是指向第l$l$层第j$j$个神经元的，不与第l$l$行第i$i$个神经元相接，因此所有偏层均为0，如下所示：

zliw(j)=[00...0]∈RNl−1(3.2.009)$$\begin{array}{}\text{(3.2.009)}& \frac{z_i^l}{{\mathit{w}}^{(j)}}=\left[\begin{array}{cccc}0& 0& ...& 0\end{array}\right]\in R^{N_{l-1}}\end{array}$$

如果i=j$i=j$时，w(j)${\mathit{w}}^{(j)}$是由指向第l$l$层第i$i$个神经元的所有连接权值组成的，根据输入信号定义可得：

zliw(j)=[∂zli∂Wlj,1∂zli∂Wlj,2...∂zli∂Wlj,Nl−1]=[al−11al−12...al−1Nl−1](3.2.010)$$\begin{array}{}\text{(3.2.010)}& \frac{z_i^l}{{\mathit{w}}^{(j)}}=\left[\begin{array}{cccc}\frac{\mathrm{\partial }{z}_i^l}{\mathrm{\partial }{W}_{j,1}^l}& \frac{\mathrm{\partial }{z}_i^l}{\mathrm{\partial }{W}_{j,2}^l}& ...& \frac{\mathrm{\partial }{z}_i^l}{\mathrm{\partial }{W}_{j,N_{l-1}}^l}\end{array}\right]=\left[\begin{array}{cccc}{a}_1^{l-1}& a_2^{l-1}& ...& a_{N_{l-1}}^{l-1}\end{array}\right]\end{array}$$

因此式(3.2.008)矩阵的每个元素为指向纸里的RNl−1$R^{N_{l-1}}$向量，当不在对角线上时，所有元素值为零，当在对角线上时，元素为第l−1$l-1$层输出值。
下面我们来看，怎样通过TensorFlow Eager Excecution API来求出这些偏导的值。

@tf.custom_gradient
def f002(W, a, b):
    def grad_fn(dy):
        ws = W.shape
        pz_pW = np.zeros((2, 2, 3))
        a1 = tf.reshape(a, [3])
        for idx in range(ws[0]):
            pz_pW[idx][idx] = a1
        diag = tf.ones([W.shape[0]])
        d_b = tf.matrix_diag(diag)
        return tf.constant(pz_pW), W, d_b
    return tf.matmul(W, a) + b, grad_fn

def test001(args={}):
    tf.enable_eager_execution()
    tfe = tf.contrib.eager

    W = tf.constant([[4.0, 5.0, 6.0],[7.0, 8.0, 9.0]])
    a = tf.reshape(tf.constant([1.0, 2.0, 3.0]), [3, 1])
    i_debug = 2
    if 1 == i_debug:
        f003(W, a)
        return 
    b = tf.reshape(tf.constant([1001.0, 1002.0]), [2, 1])
    z = f002(W, a, b)
    print('z=Wa+b={0}'.format(z))
    grad_f1 = tfe.gradients_function(f002)
    dv = grad_f1(W, a, b)
    print('pz_pW={0}'.format(dv[0].numpy()))
    print('pz_pa={0}'.format(dv[1].numpy()))
    print('pz_pb={0}'.format(dv[2].numpy()))
    print('v0.0.1')

在上面的代码中，需要说明的是第16行，将向量a$\mathit{a}$定义为3行1列的矩阵形式，这是为了与连接权值矩阵做矩阵相乘。
在前向计算阶段，可以直接调用TensorFlow的矩阵乘法和加法，我们就可以取得正确的结果，但是在反向求导阶段，如果我们直接利用TensorFlow进行求导，例如求∂zl∂al−1$\frac{\mathrm{\partial }{\mathit{z}}^l}{\mathrm{\partial }{\mathit{a}}^{l-1}}$时，TensorFlow返回结果维数与al−1${\mathit{a}}^{l-1}$相同，而Jacobian矩阵的维数为RNl×Nl−1$R^{N_l\times N_{l-1}}$，所以我们需要自己定义求导函数，根据上面的理论分析，∂zl∂Wl$\frac{\mathrm{\partial }{\mathit{z}}^l}{\mathrm{\partial }{W}^l}$是一个3维的张量，维数为RNl×Nl×Nl−1$R^{N_l\times N_l\times N_{l-1}}$，我们可以将其视为一个RNl×Nl$R^{N_l\times N_l}$的矩阵，矩阵中每个元素均为一个数组，长度为Nl−1$N_{l-1}$，且除对角线上的元素外，数组元素为0，而对角线上的元素，数组元素为第l−1$l-1$导神经元的输出值。∂zl∂al−1$\frac{\mathrm{\partial }{\mathit{z}}^l}{\mathrm{\partial }{\mathit{a}}^{l-1}}$为第l−1$l-1$层到第l$l$层的连接权值矩阵Wl$W^l$，而∂zl∂bl$\frac{\mathrm{\partial }{\mathit{z}}^l}{\mathrm{\partial }{\mathit{b}}^l}$为RNl×Nl$R^{N_l\times N_l}$的单位阵，运行结果如下所示：

至此我们完成了第l−1$l-1$层到第l$l$层正向传输和反向求导工作，基本上按照数学理论要求，我们就可以完全处理一个多层感知器模型了。