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 层做一些预处理。

来看下面这个例子：

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，无需预处理直接向下传递即可：

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 的一些基本属性。模型要使用的所有层 / 模块都需要在构造函数中声明。

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 表示的是模型实现的逻辑。来看一个例子：

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 优化器为例，优化神经网络模型参数的基本流程如下：

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，荣誉证书、定制礼品享不停。