将动态图转为静态图（Trace）#

备注

一般的模型开发阶段推荐用动态图，容易进行调试。

使用 trace 装饰器#

MegEngine 提供了很方便的动静态图转换的方法，几乎无需代码改动即可实现转换。

假设我们写好了一份动态图代码，其中训练部分代码如下：

for epoch in range(total_epochs):
    total_loss = 0
    for step, (batch_data, batch_label) in enumerate(dataloader):
        data = mge.tensor(batch_data)
        label = mge.tensor(batch_label)

        with gm:
            logits = model(data)
            loss = F.loss.cross_entropy(logits, label)
            gm.backward(loss)
            optimizer.step().clear_grad()

        total_loss += loss.numpy().item()
    print("epoch: {}, loss {}".format(epoch, total_loss/len(dataloader)))

我们可以通过以下三步将上面的动态图转换为静态图：

  1. 将循环内的前向计算、反向传播和参数优化代码提取成单独的函数，如下面例子中的 train_func() ；

  2. 将网络所需输入作为训练函数的参数，并返回任意你需要的结果（如输出结果、损失函数值等）；

  3. 用 jit 模块中的 trace 装饰器来装饰这个函数，将其中的代码变为静态图代码。

修改后的代码如下：

from megengine.jit import trace

@trace
def train_func(data, label, *, opt, gm, net):
    with gm:
        logits = model(data)
        loss = F.loss.cross_entropy(logits, label)
        gm.backward(loss)
        opt.step().clear_grad()
    return loss

for epoch in range(total_epochs):
    total_loss = 0
    for step, (batch_data, batch_label) in enumerate(dataloader):
        data = mge.tensor(batch_data)
        label = mge.tensor(batch_label)

        loss = train_func(data, label, opt=optimizer, gm=gm, net=model)
        total_loss += loss.numpy().item()
    print("epoch: {}, loss {}".format(epoch, total_loss/len(dataloader)))

对于上述代码，我们作进一步的解释：

  • jit ： 即时编译 （Just-in-time compilation）的缩写，这里作为整个静态图相关模块的名字。

  • trace ：得到静态图的一种方式，直译为 “追溯”， 含义为通过追溯输出（比如损失值、预测值）所依赖的网络结构，得到整体的计算图，再进行编译。

  • 参数列表 ： trace 在编译静态图时会根据传入参数是位置参数还是关键字参数来采取不同的处理方式。 其中位置参数用于传入网络的输入如数据和标签，关键字参数用于传入其它变量，如网络和优化器等。

trace 进阶设置#

指定静态图构造方式#

MegEngine 在编译静态图时有 “动态构造” 和 “静态构造” 两种模式（默认使用前者）。

在绝大部分情况下，两种模式下构造出的静态图并没有区别，使用中也没有分别。

我们可以指定 symbolic 参数来指定构造方式，示例代码如下:

from megengine.jit import trace

@trace(symbolic=True)
def train_func(data, label, *, opt, gm, net):
    pass

设置为 True 表示 “静态构造” 或者 “根据符号构造” ——

  • 原理： 此时，计算图中的所有数据节点（即张量）被视为符号（即 symbolic ）。 它们仅仅作为占位符（Placeholder），不产生实际的内存分配，也没有实际的值。 此时计算图的编译过程完全取决于计算图的结构，而不取决于张量的具体值，是真正的 “静态”。

  • 优点： 始终高效，能充分利用静态图的内存优化。

  • 缺点： 如果网络中包含了需要运行时动态信息才能计算的条件语句，将会失败。

设置为 False 表示 “动态构造” 或者 “根据值构造” ——

  • 原理： 被装饰的函数在第一次被调用时会根据输入的数据执行一次计算构建出一个动态图。 接着将这个动态图会被编译静态图。此后该函数的所有调用都会运行这个静态图，而不再依赖调用时输入的值。 此种模式可以视为 “动态构建第一次，此后静态运行”。

  • 优点： 根据第一次运行时信息的不同，可以构建出不同的静态图。

  • 缺点： 由于第一次的运行在动态图模式下，无法利用静态图的内存优化，通常会耗费更大的内存。 这可能导致本来在静态图模式下可以运行的网络，在第一次运行时由于内存不够而失败。

警告

在动态构造模式（设置为 False）下，如果条件语句出现在循环语句内，在循环的第一次执行中构造出的静态图将固定不再改变 （即使在循环的后续执行中，该条件语句的结果发生了变化）

将参数固定以便导出#

有的时候我们希望将一些参数（比如卷积层的卷积核等）固化下来，因此需要指定 capture_as_const = True :

from megengine.jit import trace

@trace(capture_as_const = True)
def train_func(data, label, *, opt, gm, net):
    pass

备注

如果想要使用 dump 导出模型序列化文件并进行后续处理， 则必须在 trace 时固定参数。

减少访存操作实现加速#

通常，模型中不仅含有计算受限的操作，还含有一些访存受限操作（如 Elemwsie ）. MegEngine 内嵌了 Codegen 优化机制，它可以在运行时将模型中多个操作融合起来， 并生成可以在目标机器上运行的代码，以此减少访存操作从而达到加速的目的。

备注

  • 我们在 trace 接口中传入 symbolic=True, opt_level=3, 即可打开 Codegen 优化；

  • 关于 symbolic 参数的说明，请参考 指定静态图构造方式 。

MegEngine 的 Codegen 目前集成了三种后端，分别是 NVRTC, HALIDE 和 MLIR. 其中 NVRTC 和 HALIDE 仅支持在 GPU 上使用，MLIR 则同时支持 GPU 和 CPU, 不同的后端生成代码的策略有所不同，所以运行效率也各异。 我们可以通过设置 MGB_JIT_BACKEND 环境变量来改变 Codegen 的后端，例如：

export MGB_JIT_BACKEND="NVRTC"

该环境变量在 NVIDIA GPU 环境下可取的值为 NVRTC, HALIDE 和 MLIR, 默认值为 HALIDE.

对于 CPU, 目前暂时仅支持 MLIR 后端。

警告

如果想要使用 MLIR 后端, 需要单独编译 MegEngine. 在使用 CMake 时换成如下命令：

cmake .. -DMGE_WITH_JIT=ON -DMGE_WITH_JIT_MLIR=ON -DMGE_WITH_HALIDE=OFF

然后设置如下的环境变量：

export MGB_JIT_BACKEND="MLIR"

指定代码不被转换#

使用 exclude_from_trace ，其中的代码不会被 trace, 且允许访问静态区域的 Tensor .

示例代码如下：

from megengine import jit, tensor

@jit.trace
def f(x):
    x += 1
    with jit.exclude_from_trace():  # 不对下面的 if 语句进行 trace
        if i % 2 == 0:
            x += 1
    return x

for i in range(3):
    x = tensor([1])
    print(f(x))

输出为：

Tensor([3], dtype=int32, device=xpux:0)
Tensor([2], dtype=int32, device=xpux:0)
Tensor([3], dtype=int32, device=xpux:0)