在前面的文章中，我們簡單介紹了在 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，榮譽證書、定製禮品享不停。