MegEngine Python 层模块串讲(中)

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


(资料图)

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

Python 层计算接口 —— functional 模块

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

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

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

  • builtin

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

  • apply

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

Functional中的许多op都需要通过builtinapply调用底层MegDNNop来进行计算操作。然而在实际的计算发生前,很多时候需要在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数量、devicepython层做了一些预处理,然后才调用builtinapply向下转发。

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

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

对于实现了对应kernelop,其在imperative层的实现通常非常的短。

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

部分functionalop不直接调用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]:返回一个此moduleTensor的可迭代对象;
  • def children(self, **kwargs) -> "Iterable[Module]":返回一个可迭代对象,该对象包括属于当前模块的直接属性的子模块;
  • def named_buffers(self, prefix: Optional[str] = None, recursive: bool = True, **kwargs) -> Iterable[Tuple[str, Tensor]]:返回当前模块中keybuffer的键值对的可迭代对象,这里key是从该模块至buffer的点路径(dotted path);
  • def named_parameters(self, prefix: Optional[str] = None, recursive: bool = True, **kwargs) -> Iterable[Tuple[str, Parameter]]:返回当前模块中keyparameter的键值对的可迭代对象,这里key是从该模块至buffer的点路径(dotted path);
  • def named_tensors(self, prefix: Optional[str] = None, recursive: bool = True, **kwargs) -> Iterable[Tuple[str, Tensor]]:返回当前模块中keyTensorbuffer + 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):返回模块的状态字典,状态字典是一个保存当前模块所有可学习的Tensorbuffer + parameter)的字典。出于兼容性考虑,字典中的value的数据结构类型为numpy.ndarray(而不是Tensor),并且不可修改,是只读的;
  • def load_state_dict(self, state_dict: Union[dict, Callable[[str, Tensor], Optional[np.ndarray]]], strict=True, ):加载一个模块的状态字典,这个方法常用于模型训练过程的保存与加载。

值得一提的是,ParametersBuffer都是与Module相关的Tensor,它们的区别可以理解为:

  • Parameter是模型的参数,在训练过程中会通过反向传播进行更新,因此值是可能改变的,常见的有weightbias等;
  • Buffer是模型用到的统计量,不会在反向传播过程中更新,常见的有meanvar等。

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

使用 optimizer 模块优化模型参数

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

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

from megengine.autodiff import GradManagerimport 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,荣誉证书、定制礼品享不停。

关键词: