引言

本着“凡我不能创造的，我就不能理解”的思想，本系列文章会基于纯Python以及NumPy从零创建自己的深度学习框架，该框架类似PyTorch能实现自动求导。

要深入理解深度学习，从零开始创建的经验非常重要，从自己可以理解的角度出发，尽量不适用外部完备的框架前提下，实现我们想要的模型。本系列文章的宗旨就是通过这样的过程，让大家切实掌握深度学习底层实现，而不是仅做一个调包侠。

在上篇文章中，我们实现了常见运算的计算题，本文来实现剩下的：Max、Slice、Reshape和Transpose的计算图。

求最大值

还是先写测试用例：

from core.tensor import Tensor
import numpy as np


def test_simple_max():
    x = Tensor([1, 2, 3, 6, 7, 9, 2], requires_grad=True)
    z = x.max()

    assert z.data == [9]
    z.backward()

    assert x.grad.data.tolist() == [0, 0, 0, 0, 0, 1, 0]


def test_simple_max2():
    x = Tensor([1, 2, 3, 9, 7, 9, 2], requires_grad=True)
    z = x.max()

    assert z.data == [9]  # 最大值还是9
    z.backward()

    # 但是有两个最大值，所以梯度被均分了
    assert x.grad.data.tolist() == [0, 0, 0, 0.5, 0, 0.5, 0]


def test_matrix_max():
    a = np.array([[1., 1., 8., 9., 1.],
                  [4., 5., 9., 9., 8.],
                  [8., 6., 9., 7., 9.],
                  [8., 6., 1., 9., 8.]])

    x = Tensor(a, requires_grad=True)
    z = x.max()

    assert z.data == [9]  # 最大值是9
    z.backward()

    # 总共有6个9
    np.testing.assert_array_almost_equal(x.grad.data, [[0, 0, 0, 1 / 6, 0],
                                                       [0, 0, 1 / 6, 1 / 6, 0],
                                                       [0, 0, 1 / 6, 0, 1 / 6],
                                                       [0, 0, 0, 1 / 6, 0]])


def test_matrix_max2():
    a = np.array([[1., 1., 8., 9., 1.],
                  [4., 5., 9., 9., 8.],
                  [8., 6., 9., 7., 9.],
                  [8., 6., 1., 9., 8.]])

    x = Tensor(a, requires_grad=True)
    z = x.max(axis=0)  # [8, 6, 9, 9, 9]

    assert z.data.tolist() == [8, 6, 9, 9, 9]
    z.backward([1, 1, 1, 1, 1])

    grad = [[0., 0., 0., 1 / 3, 0.],
            [0., 0., 0.5, 1 / 3, 0.],
            [0.5, 0.5, 0.5, 0, 1],
            [0.5, 0.5, 0., 1 / 3, 0.]]

    np.testing.assert_array_almost_equal(x.grad.data, np.array(grad))

分析的代码在文章计算图运算补充中，这里就不再赘述。

class Max(_Function):
    def forward(ctx, x: ndarray, axis=None, keepdims=False) -> ndarray:
        ret = np.amax(x, axis=axis, keepdims=keepdims)
        ctx.save_for_backward(x, axis, ret, keepdims)
        return ret

    def backward(ctx, grad: ndarray) -> ndarray:
        x, axis, ret, keepdims = ctx.saved_tensors
        mask = (x == ret)
        div = mask.sum(axis=axis, keepdims=keepdims)
        return mask * grad / div

切片

切片就是索引操作，测试代码如下：

from core.tensor import Tensor
import numpy as np


def test_get_by_index():
    x = Tensor([1, 2, 3, 4, 5, 6, 7], requires_grad=True)
    z = x[2]

    assert z.data == 3
    z.backward()

    assert x.grad.data.tolist() == [0, 0, 1, 0, 0, 0, 0]


def test_slice():
    x = Tensor([1, 2, 3, 4, 5, 6, 7], requires_grad=True)
    z = x[2:4]

    assert z.data.tolist() == [3, 4]
    z.backward([1, 1])

    assert x.grad.data.tolist() == [0, 0, 1, 1, 0, 0, 0]


def test_matrix_slice():
    a = np.array([[1., 1., 8., 9., 1.],
                  [4., 5., 9., 9., 8.],
                  [8., 6., 9., 7., 9.],
                  [8., 6., 1., 9., 8.]])

    x = Tensor(a, requires_grad=True)
    z = x[1:3, 2:4]

    assert z.data.tolist() == [[9, 9], [9, 7]]
    z.backward([[1, 1], [1, 1]])

    # 总共有6个9
    np.testing.assert_array_almost_equal(x.grad.data, [[0, 0, 0, 0, 0],
                                                       [0, 0, 1, 1, 0],
                                                       [0, 0, 1, 1, 0],
                                                       [0, 0, 0, 0, 0]])

class Slice(_Function):
    def forward(ctx, x: ndarray, idxs: slice) -> ndarray:
        '''
        z = x[idxs]
        '''
        # 如果传入[1:3]，变成切片slice
        # 如果idxs传入单个索引，会被看成是整数，所以这里转换回来
        if isinstance(idxs, ndarray):
            idxs = int(idxs.item())
        ctx.save_for_backward(x.shape, idxs)
        return x[idxs]

    def backward(ctx, grad) -> Tuple[ndarray, None]:
        x_shape, idxs = ctx.saved_tensors
        bigger_grad = np.zeros(x_shape, dtype=grad.dtype)
        bigger_grad[idxs] = grad

        return bigger_grad, None

变形

变形(Reshape)操作的反向传播其实是最简单的。假设经过y = x.reshape(..)，在反向传播时，只要保证梯度的形状和x保持一致即可。

测试用例：

import numpy as np

from core.tensor import Tensor


def test_reshape():
    x = Tensor(np.arange(9), requires_grad=True)
    z = x.reshape((3, 3))
    z.backward(np.ones((3, 3)))

    assert x.grad.data.tolist() == np.ones_like(x.data).tolist()


def test_matrix_reshape():
    x = Tensor(np.arange(12).reshape(2, 6), requires_grad=True)
    z = x.reshape((4, 3))

    z.backward(np.ones((4, 3)))

    assert x.grad.data.tolist() == np.ones_like(x.data).tolist()

代码实现：

class Reshape(_Function):
    def forward(ctx, x: ndarray, shape: Tuple) -> ndarray:
        ctx.save_for_backward(x.shape)
        return x.reshape(shape)

    def backward(ctx, grad: ndarray) -> Tuple[ndarray, None]:
        x_shape, = ctx.saved_tensors
        return grad.reshape(x_shape), None

转置

变形就是Reshape操作，在计算图运算补充中中，我们详细分析了变形和转置的区别。

比如

经过变形后：

转置：

我们实现测试用例：

import numpy as np

from core.tensor import Tensor


def test_transpose():
    x = Tensor(np.arange(6).reshape((2, 3)), requires_grad=True)
    z = x.T

    assert z.data.shape == (3, 2)
    z.backward(np.ones((3, 2)))

    assert x.grad.data.tolist() == np.ones_like(x.data).tolist()


def test_matrix_transpose():
    x = Tensor(np.arange(12).reshape((2, 6, 1)), requires_grad=True)
    z = x.transpose((0, 2, 1))

    assert z.data.shape == (2, 1, 6)

    z.backward(np.ones((2, 1, 6)))

    assert x.grad.data.tolist() == np.ones_like(x.data).tolist()

代码实现：

class Transpose(_Function):
    def forward(ctx, x: ndarray, axes) -> ndarray:
        ctx.save_for_backward(axes)
        return x.transpose(axes)

    def backward(ctx, grad: ndarray) -> Any:
        axes, = ctx.saved_tensors
        if axes is None:
            return grad.transpose()
        return grad.transpose(tuple(np.argsort(axes))), None

完整代码

完整代码笔者上传到了程序员最大交友网站上去了，地址: 👉  https://github.com/nlp-greyfoss/metagrad

总结

到此，基本上我们会用到基本运算的计算图都实现了。从下篇文章开始，就基于我们的自动求导工具来实现深度学习模型了。