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 GradManagerimport megengine.optimizer as optimmodel = MyModel()gm = GradManager().attach(model.parameters())optimizer = optim.SGD(model.parameters(), lr=0.01) # lr may vary with different modelfor 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,荣誉证书、定制礼品享不停。