从零实现深度学习框架——优化索引操作&交叉熵损失

引言

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

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

为了继续下去，之前的某些实现已经不能满足现在的需要了。

比如交叉熵损失函数只支持one-hot作为target，这在one-hot的维度很大时，非常低效。

现在按照这个思路，首先为​​Tensor​​​添加数据类型​​dtype​​​，然后优化索引实现，让其支持​​Integer array indexing​​，最后优化交叉熵损失函数，让其同时支持one-hot和类别索引作为target。

支持指定数据类型

首先在初始化​​Tensor​​时可以指定数据类型：

def __init__(self, data: Arrayable, requires_grad: bool = False, dtype=None) -> None:
    '''
        初始化Tensor对象
        Args:
            data: 数据
            requires_grad: 是否需要计算梯度
            dtype: 数据类型，默认为None
        '''

默认不传时，会根据实际类型判断数据类型：

def ensure_array(arrayable: Arrayable, dtype=None) -> np.ndarray:
    """
    :param arrayable:
    :return:
    """
    if isinstance(arrayable, Number):
        if dtype is None:
            dtype = type(arrayable)
        return np.array(arrayable, dtype=dtype)
    elif isinstance(arrayable, list):
        # 让np自己判断数据类型
        return np.array(arrayable, dtype=dtype)
    else:
        return

对于​​Number​​​和​​list​​，我们转换为Numpy数组，否则不处理。

最后，对一些静态初始化方法增加数据类型支持：

@classmethod
    def empty(cls, *shape, dtype=_type, **kwargs):
        return cls(np.empty(*shape, dtype=dtype), **kwargs)

    @classmethod
    def zeros(cls, *shape, dtype=_type, **kwargs) -> "Tensor":
        return cls(np.zeros(shape, dtype=dtype), **kwargs)

    @classmethod
    def ones(cls, *shape, dtype=_type, **kwargs) -> "Tensor":
        return cls(np.ones(shape, dtype=dtype), **kwargs)

    @classmethod
    def ones_like(cls, t: "Tensor", dtype=_type, **kwargs) -> "Tensor":
        return cls(np.ones(t.shape, dtype=dtype), **kwargs)

优化索引实现

有了上面的基础，实现索引就很容易了。

class Slice(Function):
    def forward(ctx, x: ndarray, slices: Any) -> ndarray:
        '''
        z = x[slices]
        '''
        ctx.save_for_backward(x.shape, slices)
        return x[slices]

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

        return bigger_grad, None

可以充分利用Numpy实现我们想要的功能：

def test_boolean_indexing():
    '''测试boolean索引操作'''
    x = Tensor([1, 2, 3, 4, 5, 6, 7], requires_grad=True)
    z = x[x < 5]

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

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


def test_integer_indexing():
    x = Tensor(np.arange(35).reshape(5, 7), requires_grad=True)
    # Tensor([[(0)  1  2  3  4  5  6]
    #         [7  8  9 10 11 12 13]
    #         [14 (15) 16 17 18 19 20]
    #         [21 22 23 24 25 26 27]
    #         [28 29 (30) 31 32 33 34]], requires_grad = True)
    #

    # ! z = x[Tensor([0, 2, 4]), Tensor([0, 1, 2])] 暂不支持元组Tensor作为索引
    z = x[np.array([0, 2, 4]), np.array([0, 1, 2])]  # x[0,0] x[2,1] x[4,2]

    assert z.data.tolist() == [0, 15, 30]

    z.sum().backward()

    assert x.grad.data.tolist() == [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                                    [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                                    [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0]]

现在也支持​​x[np.array([0, 2, 4]), np.array([0, 1, 2])]​​

即获取 ​​x[0,0] x[2,1] x[4,2]​​

我们先来优化​​NLLLoss​​的实现，然后基于此实现交叉熵损失。

优化负对数似然损失

负对数似然损失的公式在Softmax回归中的数值稳定​有过详细介绍，这里我们回顾下它的公式：

这里是logits，是真实类别。我们现在支持损失中传入的​​​target​​为类别索引或one-hot向量：

def nll_loss(input: Tensor, target: Tensor, reduction: str = "mean") -> Tensor:
    '''
    负对数似然损失
    :param input: 对数概率 即 log_softmax
    :param target:  类别索引 或 one-hot向量
    :param reduction:
    :return:
    '''
    # 如果target是ont-hot向量
    if input.ndim == target.ndim and input.shape == target.shape:
        errors = - target * input
    else:
        # 如果target是类别索引
        errors = -input[range(target.shape[0]), target.numpy()]
    return _reduction(errors, reduction)

这里​​input​​要求输入的就是概率的对数，即logsoftmax。

如果​​target​​是类别索引的话，那么可以通过numpy整数数组索引快速取值。

假设有两个样本，属于四个类别，经过Softmax得到的概率为：

模型判断​​样本1​​​属于​​类别4​​​；判断​​样本2​​​属于​​类别1​​​。那么取对数之后​​input​​变成：

假设此时输入的类别索引分别为​​3​​​和​​0​​​。即第一个样本的真实标签为​​类别4​​​;第二个样本的真实标签为​​类别1​​。

那么​​-input[range(2),numpy([3,0]) ]​​​的结果为取​​input[0,3]​​​的和​​input[1,0]​​​，即选取了样本真实标签对应的输出值。正如公式所表达的。

负对数似然实现起来真的很简单，尤其是传入的类别索引。

优化交叉熵损失函数

我们这次基于负对数似然来实现交叉熵损失。那么实现起来不要太简单：

def cross_entropy(input: Tensor, target: Tensor, reduction: str = "mean") -> Tensor:
    '''

    :param input: logits
    :param target: 真实标签one-hot向量 或 类别索引
    :param reduction:
    :return:
    '''
    # 先计算logsoftmax
    log_y = log_softmax(input)
    # 基于nll实现交叉熵损失
    return nll_loss(log_y, target, reduction)

完整代码

点击 👉 ​​https://github.com/nlp-greyfoss/metagrad​​