昇思MindSpore进阶教程-Cell

大家好，我是刘明，明志科技创始人，华为昇思MindSpore布道师。
技术上主攻前端开发、鸿蒙开发和AI算法研究。
努力为大家带来持续的技术分享，如果你也喜欢我的文章，就点个关注吧

Cell作为神经网络构造的基础单元，与神经网络层(Layer)的概念相对应，对Tensor计算操作的抽象封装，能够更准确清晰地对神经网络结构进行表示。除了基础的Tensor计算流程定义外，神经网络层还包含了参数管理、状态管理等功能。

Cell训练状态转换

神经网络中的部分Tensor操作在训练和推理时的表现并不相同，如nn.Dropout在训练时进行随机丢弃，但在推理时则不丢弃，nn.BatchNorm在训练时需要更新mean和var两个变量，在推理时则固定其值不变。因此我们可以通过Cell.set_train接口来设置神经网络的状态。

set_train(True)时，神经网络状态为train, set_train接口默认值为True:

net.set_train()
print(net.phase)

set_train(False)时，神经网络状态为predict：

net.set_train(False)
print(net.phase)

自定义神经网络层

通常情况下，MindSpore提供的神经网络层接口和function函数接口能够满足模型构造需求，但由于AI领域不断推陈出新，因此有可能遇到新网络结构没有内置模块的情况。此时我们可以根据需要，通过MindSpore提供的function接口、Primitive算子自定义神经网络层，并可以使用Cell.bprop方法自定义反向。下面分别详述三种自定义方法。

使用function接口构造神经网络层

MindSpore提供大量基础的function接口，可以使用其构造复杂的Tensor操作，封装为神经网络层。

自定义Cell反向

在特殊场景下，我们不但需要自定义神经网络层的正向逻辑，也需要手动控制其反向的计算，此时我们可以通过Cell.bprop接口对其反向进行定义。在全新的神经网络结构设计、反向传播速度优化等场景下会用到该功能。下面我们以Dropout2d为例，介绍如何自定义Cell反向：

class Dropout2d(nn.Cell):
    def __init__(self, keep_prob):
        super().__init__()
        self.keep_prob = keep_prob
        self.dropout2d = ops.Dropout2D(keep_prob)

    def construct(self, x):
        return self.dropout2d(x)

    def bprop(self, x, out, dout):
        _, mask = out
        dy, _ = dout
        if self.keep_prob != 0:
            dy = dy * (1 / self.keep_prob)
        dy = mask.astype(mindspore.float32) * dy
        return (dy.astype(x.dtype), )

dropout_2d = Dropout2d(0.8)
dropout_2d.bprop_debug = True

bprop方法分别有三个入参：

• x: 正向输入，当正向输入为多个时，需同样数量的入参。
• out: 正向输出。
• dout: 反向传播时，当前Cell执行之前的反向结果。
  一般我们需要根据正向输出和前层反向结果配合，根据反向求导公式计算反向结果，并将其返回。Dropout2d的反向计算需要根据正向输出的mask矩阵对前层反向结果进行mask，然后根据keep_prob进行缩放。最终可得到正确的计算结果。

自定义Cell反向时，在PyNative模式下支持拓展写法，可以对Cell内部的权重求导，具体列子如下：

class NetWithParam(nn.Cell):
    def __init__(self):
        super(NetWithParam, self).__init__()
        self.w = Parameter(Tensor(np.array([2.0], dtype=np.float32)), name='weight')
        self.internal_params = [self.w]

    def construct(self, x):
        output = self.w * x
        return output

    def bprop(self, *args):
        return (self.w * args[-1],), {self.w: args[0] * args[-1]}

bprop方法支持*args入参，args数组中最后一位args[-1]为返回给该cell的梯度。通过self.internal_params设置求导的权重，同时在bprop函数的返回值为一个元组和一个字典，返回输入对应梯度的元组，以及以key为权重，value为权重对应梯度的字典。

Hook功能

调试深度学习网络是每一个深度学习领域的从业者需要面对且投入精力较大的工作。由于深度学习网络隐藏了中间层算子的输入、输出数据以及反向梯度，只提供网络输入数据（特征量、权重）的梯度，导致无法准确地感知中间层算子的数据变化，从而降低了调试效率。为了方便用户准确、快速地对深度学习网络进行调试，MindSpore在动态图模式下设计了Hook功能，使用Hook功能可以捕获中间层算子的输入、输出数据以及反向梯度。

目前，动态图模式下提供了四种形式的Hook功能，分别是：HookBackward算子和在Cell对象上进行注册的register_forward_pre_hook、register_forward_hook、register_backward_hook功能。

HookBackward算子

HookBackward将Hook功能以算子的形式实现。用户初始化一个HookBackward算子，将其安插到深度学习网络中需要捕获梯度的位置。在网络正向执行时，HookBackward算子将输入数据不做任何修改后原样输出；在网络反向传播梯度时，在HookBackward上注册的Hook函数将会捕获反向传播至此的梯度。用户可以在Hook函数中自定义对梯度的操作，比如打印梯度，或者返回新的梯度。

Cell对象的register_forward_pre_hook功能

用户可以在Cell对象上使用register_forward_pre_hook函数来注册一个自定义的Hook函数，用来捕获正向传入该Cell对象的数据。该功能在静态图模式下和在使用@jit修饰的函数内不起作用。register_forward_pre_hook函数接收Hook函数作为入参，并返回一个与Hook函数一一对应的handle对象。用户可以通过调用handle对象的remove()函数来删除与之对应的Hook函数。每一次调用register_forward_pre_hook函数，都会返回一个不同的handle对象。

Cell对象的register_forward_hook功能

用户可以在Cell对象上使用register_forward_hook函数来注册一个自定义的Hook函数，用来捕获正向传入Cell对象的数据和Cell对象的输出数据。该功能在静态图模式下和在使用@jit修饰的函数内不起作用。register_forward_hook函数接收Hook函数作为入参，并返回一个与Hook函数一一对应的handle对象。用户可以通过调用handle对象的remove()函数来删除与之对应的Hook函数。每一次调用register_forward_hook函数，都会返回一个不同的handle对象。

Cell对象的register_backward_hook功能

用户可以在Cell对象上使用register_backward_hook函数来注册一个自定义的Hook函数，用来捕获网络反向传播时与Cell对象相关联的梯度。该功能在图模式下或者在使用@jit修饰的函数内不起作用。register_backward_hook函数接收Hook函数作为入参，并返回一个与Hook函数一一对应的handle对象。用户可以通过调用handle对象的remove()函数来删除与之对应的Hook函数。每一次调用register_backward_hook函数，都会返回一个不同的handle对象。

与HookBackward算子所使用的自定义Hook函数有所不同，register_backward_hook使用的Hook函数的入参中，包含了表示Cell对象名称与id信息的cell_id、反向传入到Cell对象的梯度、以及Cell对象的反向输出的梯度。