backprop_v2.py

import autograd.numpy as np
from autograd import grad


def main():

    def relu(x):
        return np.maximum(x, np.zeros_like(x))

    def drelu(x):
        ret = np.zeros_like(x)
        ret[x > 0] = 1.0
        return ret

    d = 5
    n1 = 10
    n2 = 15

    xpt = np.random.randn(d)
    ypt = np.random.randn(n2)

    def net(args):
        W1, W2, b1 = args
        z = np.dot(W2, relu(np.dot(W1, xpt) + b1)) - ypt
        return 0.5 * np.dot(z, z)

    g = grad(net)

    W1 = np.random.randn(n1, d)
    W2 = np.random.randn(n2, n1)
    b1 = np.random.randn(n1)

    ret1 = g((W1, W2, b1))

    def compute_grad(args):
        W1, W2, b1 = args

        z1 = np.dot(W1, xpt) + b1
        z2 = relu(z1)
        z3 = np.dot(W2, z2)
        z4 = z3 - ypt
        f = 0.5 * np.dot(z4, z4)

        Df_z4 = z4.reshape((1, n2))
        Df_z3 = Df_z4
        Df_z2 = np.dot(Df_z3, W2)
        Df_z1 = np.dot(Df_z2, np.diag(drelu(z1)))
        Df_b1 = Df_z1

        return (np.outer(Df_z1.flatten(), xpt),
                np.outer(Df_z3.flatten(), z2),
                Df_b1.flatten())

    ret2 = compute_grad((W1, W2, b1))

    assert np.allclose(ret1[0], ret2[0])
    assert np.allclose(ret1[1], ret2[1])
    assert np.allclose(ret1[2], ret2[2])


if __name__ == '__main__':
    main()