MegEngine Python 层模块串讲（中）

在前面的文章中，我们简单介绍了在 MegEngine imperative 中的各模块以及它们的作用。对于新用户而言可能不太了解各个模块的使用方法，对于模块的结构和原理也是一头雾水。Python 作为现在深度学习领域的主流编程语言，其相关的模块自然也是深度学习框架的重中之重。

模块串讲将对 MegEngine 的 python 层相关模块分别进行更加深入的介绍，会涉及到一些原理的解释和代码解读。Python 层模块串讲共分为上、中、下三个部分，本文将介绍 Python 层的 functional、module 和 optimizer 模块。理解并掌握这几个模块对于高效搭建神经网络非常重要。

Python 层计算接口 ------ functional 模块

我们在定义网络结构时经常需要包含一些计算操作，这些计算操作就定义在 functional 中。

functional 中实现了各类计算函数，包含对很多 op 的封装，供实现模型时调用。

functional 中有些 op 完全是由 Python 代码实现，有些则需要调用 C++ 接口完成计算（没错，这里的计算就需要 MegDNN kernel）。对于后者，需要有机制确保我们的实现能够转发到底层正确执行，所以你在 functional 的许多 op 实现中会看到 builtin 和 apply：

builtin

builtin 封装了所有的 op，我们在 functional 中通过 builtin.SomeOp(param) 的方式获得一个算子 SomeOp，param 表示获取 SomeOp 需要的参数。
apply

通过 builtin 获取到 op 后，需要调用 apply 接口来调用 op 的底层实现进行计算。apply 是在 Python 层调用底层 op 的接口，apply 的第一个参数是 op（通过 builtin 获得），后面的参数是执行这个 op 需要的参数，都为 Tensor。在 imperative 中 op 计算通过 apply(op, inputs) 的方式向下传递并最终调用到MegDNN中的 kernel``。

Functional 中的许多 op 都需要通过 builtin 和 apply 调用底层 MegDNN 的 op 来进行计算操作。然而在实际的计算发生前，很多时候需要在 Python 层做一些预处理。

来看下面这个例子：

C++ 复制代码

def concat(inps: Iterable[Tensor], axis: int = 0, device=None) -> Tensor:
    ...
    if len(inps) == 1:
        return inps[0]

    if device is None:
        device = get_device(inps)
    device = as_device(device)
    (result,) = apply(builtin.Concat(axis=axis, comp_node=device.to_c()), *inps)
    return result

这里 concat 方法先对输入 tensor 数量、device 在 python 层做了一些预处理，然后才调用 builtin 和 apply 向下转发。

而对于 diag 这个 op，无需预处理直接向下传递即可：

C++ 复制代码

def diag(inp, k=0) -> Tensor:
    ...
    op = builtin.Diag(k=k)
    (result,) = apply(op, inp)
    return result

对于实现了对应 kernel 的 op，其在 imperative 层的实现通常非常的短。

上面 concat 和 diag 的 apply 调用会进入 py_apply 函数，并通过解析 Python 中的参数，将它们转换成 C++ 中的对应类型，然后调用 imperative::apply，进入 dispatch 层。

部分 functional 的 op 不直接调用 py_apply 而是有对应的 cpp 实现，比如 squeeze：

ç++ 复制代码

def squeeze(inp: Tensor, axis: Optional[Union[int, Sequence[int]]] = None) -> Tensor:
    return squeeze_cpp(inp, axis)

这样的实现往往是需要在调用 py_apply 之前做一些预处理，但使用 python 实现性能较差，所以我们将相关预处理以及 py_apply 的逻辑在 C++ 层面实现。

本文主要介绍 Python 层的方法，关于 C++ 部分的实现会在之后的文章进行更深入的介绍。

在这里我们只需要知道，functional 中包装了所有关于 Tensor 计算相关的接口，是所有计算的入口，实际的计算操作通常会被转发到更底层的 C++ 实现。

用户可以参考官方文档获取所有 functional 中的方法介绍。

模块结构的小型封装版本 ------ module 模块

神经网络模型是由对输入数据执行操作的各种层（Layer），或者说模块（Module）组成。

Module 用来定义网络模型结构，用户实现算法时要用组合模块 Module (megengine/module) 的方式搭建模型，定义神经网络时有些结构经常在模型中反复使用，将这样的结构封装为一个 Module，既可以减少重复代码也降低了复杂模型编码的难度。

一个 module 类主要有两类函数：

__init__：构造函数，定义了模型各个层的大小。用户自定义的 Module 都源自基类 class Module，所以在构造函数中一定要先调用 super().__init__()，设置 Module 的一些基本属性。模型要使用的所有层 / 模块都需要在构造函数中声明。

C++ 复制代码

class Module(metaclass=ABCMeta):
    r"""Base Module class.

    Args:
        name: module's name, can be initialized by the ``kwargs`` parameter
            of child class.
    """

    def __init__(self, name=None):
        self._modules = []

        if name is not None:
            assert (
                isinstance(name, str) and name.strip()
            ), "Module's name must be a non-empty string"

        self.name = name

        # runtime attributes
        self.training = True
        self.quantize_disabled = False

        # hooks
        self._forward_pre_hooks = OrderedDict()
        self._forward_hooks = OrderedDict()

        # used for profiler and automatic naming
        self._name = None
        self._short_name = None

    # 抽象方法，由继承的 Module 自己实现
    @abstractmethod
    def forward(self, inputs):
        pass
    
    # 其他方法
    ...

forward：定义模型结构，实现前向传播，也就是将数据输入模型到输出的过程。这里会调用 Functional (megengine/functional) 中的函数进行前向计算，forward 表示的是模型实现的逻辑。来看一个例子：

C++ 复制代码

class Simple(Module):
    def __init__(self):
        super().__init__()
        self.a = Parameter([1.23], dtype=np.float32)

    def forward(self, x):
        x = x * self.a
        return x

__init__ 表明模型中有一个参数 a，它的初值是固定的，forward 中实现了具体的计算逻辑，也就是对传入的参数与 a 进行乘法运算。

对于一些更复杂的计算操作（如卷积、池化等）就需要借助 functional 中提供的方法来完成。

除了 __init__ 和 forward，基类 class Module 提供了很多属性和方法，常用的有：

def buffers(self, recursive: bool = True, **kwargs) -> Iterable[Tensor]：返回一个可迭代对象，遍历当前模块的所有 buffers；
def parameters(self, recursive: bool = True, **kwargs) -> Iterable[Parameter]：返回一个可迭代对象，遍历当前模块所有的 parameters；
def tensors(self, recursive: bool = True, **kwargs) -> Iterable[Parameter]：返回一个此 module 的 Tensor 的可迭代对象；
def children(self, **kwargs) -> "Iterable[Module]"：返回一个可迭代对象，该对象包括属于当前模块的直接属性的子模块；
def named_buffers(self, prefix: Optional[str] = None, recursive: bool = True, **kwargs) -> Iterable[Tuple[str, Tensor]]：返回当前模块中 key 与 buffer 的键值对的可迭代对象，这里 key 是从该模块至 buffer 的点路径（dotted path）；
def named_parameters(self, prefix: Optional[str] = None, recursive: bool = True, **kwargs) -> Iterable[Tuple[str, Parameter]]：返回当前模块中 key 与 parameter 的键值对的可迭代对象，这里 key 是从该模块至 buffer 的点路径（dotted path）；
def named_tensors(self, prefix: Optional[str] = None, recursive: bool = True, **kwargs) -> Iterable[Tuple[str, Tensor]]：返回当前模块中 key 与 Tensor（buffer + parameter）的键值对的可迭代对象，这里 key 是从该模块至 Tensor 的点路径（dotted path）；
def named_modules(self, prefix: Optional[str] = None, **kwargs) -> "Iterable[Tuple[str, Module]]"：返回一个可迭代对象，该对象包括当前模块自身在内的其内部所有模块组成的 key-module 键-模块对，这里 key 是从该模块至各子模块的点路径（dotted path）；
def named_children(self, **kwargs) -> "Iterable[Tuple[str, Module]]"：返回一个可迭代对象，该对象包括当前模块的所有子模块（submodule）与键（key）组成的 key-submodule 对，这里 key 是子模块对应的属性名；
def state_dict(self, rst=None, prefix="", keep_var=False)：返回模块的状态字典，状态字典是一个保存当前模块所有可学习的 Tensor （buffer + parameter）的字典。出于兼容性考虑，字典中的 value 的数据结构类型为 numpy.ndarray （而不是 Tensor），并且不可修改，是只读的；
def load_state_dict(self, state_dict: Union[dict, Callable[[str, Tensor], Optional[np.ndarray]]], strict=True, )：加载一个模块的状态字典，这个方法常用于模型训练过程的保存与加载。

值得一提的是，Parameters 和 Buffer 都是与 Module 相关的 Tensor，它们的区别可以理解为：

Parameter 是模型的参数，在训练过程中会通过反向传播进行更新，因此值是可能改变的，常见的有 weight、bias 等；
Buffer 是模型用到的统计量，不会在反向传播过程中更新，常见的有 mean、var 等。

在 MegEngine 的 module目录下可以看到已经有很多常见的 module 实现，用户实现自己的模型可以根据需要复用其中的模块。

使用 optimizer 模块优化模型参数

MegEngine 中的 optimizer 模块实现了基于各种常见优化策略的优化器，为用户提供了包括 SGD、ADAM 在内的常见优化器实现。这些优化器能够基于参数的梯度信息，按照算法所定义的策略执行更新。

大部分情况下用户不会自己实现优化器，这里以 SGD 优化器为例，优化神经网络模型参数的基本流程如下：

C++ 复制代码

from megengine.autodiff import GradManager
import megengine.optimizer as optim

model = MyModel()
gm = GradManager().attach(model.parameters())
optimizer = optim.SGD(model.parameters(), lr=0.01)  # lr may vary with different model

for data, label in dataset:
    with gm:
        pred = model(data)
        loss = loss_fn(pred, label)
        gm.backward()
        optimizer.step().clear_grad()

这里我们构造了一个优化器 optimizer，传入参数是 model 需要被优化的 Parameter，和 learning rate；
优化器通过执行 step() 方法进行一次优化；
优化器通过执行 clear_grad() 方法清空参数梯度。
- 为何要手动清空梯度？
  
  梯度管理器执行 backward() 方法时，会将当前计算所得到的梯度以累加的形式积累到原有梯度上，而不是直接做替换。因此对于新一轮的梯度计算，通常需要将上一轮计算得到的梯度信息清空。何时进行梯度清空是由人为控制的，这样可允许灵活进行梯度的累积。

用户也可以继承 class Optimizer，实现自己的优化器。

以上就是关于 functional，Module，optimizer 的模块的基本介绍，这几个模块是我们搭建模型训练的最核心的部分，熟悉这部分后，我们就可以高效搭建神经网络了。

附

更多 MegEngine 信息获取，您可以：查看文档和 GitHub 项目，或加入 MegEngine 用户交流 QQ 群：1029741705。欢迎参与 MegEngine 社区贡献，成为 Awesome MegEngineer，荣誉证书、定制礼品享不停。