megengine.module 模块¶

megengine.module.activation¶

class megengine.module.activation.LeakyReLU(negative_slope=0.01)[源代码]¶

基类：megengine.module.module.Module

对每个元素应用函数：

\[\text{LeakyReLU}(x) = \max(0,x) + negative\_slope \times \min(0,x)\]

或者

\[\begin{split}\text{LeakyReLU}(x) = \begin{cases} x, & \text{ if } x \geq 0 \\ negative\_slope \times x, & \text{ otherwise } \end{cases}\end{split}\]

示例代码：

import numpy as np
import megengine as mge
import megengine.module as M
data = mge.tensor(np.array([-8, -12, 6, 10]).astype(np.float32))

leakyrelu = M.LeakyReLU(0.01)
output = leakyrelu(data)
print(output.numpy())

输出：

[-0.08 -0.12  6.   10.  ]

forward(inputs)[源代码]¶

class megengine.module.activation.PReLU(num_parameters=1, init=0.25)[源代码]¶

基类：megengine.module.module.Module

对每个元素应用函数：

\[\text{PReLU}(x) = \max(0,x) + a * \min(0,x)\]

或者

\[\begin{split}\text{PReLU}(x) = \begin{cases} x, & \text{ if } x \geq 0 \\ ax, & \text{ otherwise } \end{cases}\end{split}\]

这里的 \(a\) 是一个可学习参数。当以无参数方式调用 PReLU() 时，它会在所有输入通道上使用同一个 \(a\) 参数。若以 PReLU(num_of_channels) 方式调用，每个输入通道都使用不同的 \(a\)。

参数

num_parameters (int) – 待学习的 \(a\) 参数数目，仅允许两种合法值： 1，或者输入数据通道数。默认： 1
init (float) – \(a\) 初始值。默认：0.25

示例代码：

import numpy as np
import megengine as mge
import megengine.module as M
data = mge.tensor(np.array([-1.2, -3.7, 2.7]).astype(np.float32))
prelu = M.PReLU()
output = prelu(data)
print(output.numpy())

输出：

[-0.3   -0.925  2.7  ]

forward(inputs)[源代码]¶

class megengine.module.activation.ReLU[源代码]¶

基类：megengine.module.module.Module

对每个元素应用函数：

\[\text{ReLU}(x) = \max(x, 0)\]

示例代码：

import numpy as np
import megengine as mge
import megengine.module as M
data = mge.tensor(np.array([-2,-1,0,1,2,]).astype(np.float32))
relu = M.ReLU()
output = relu(data)
with np.printoptions(precision=6):
    print(output.numpy())

输出：

[0. 0. 0. 1. 2.]

forward(x)[源代码]¶

class megengine.module.activation.Sigmoid[源代码]¶

基类：megengine.module.module.Module

对每个元素应用函数：

\[\text{Sigmoid}(x) = \frac{1}{1 + \exp(-x)}\]

示例代码：

import numpy as np
import megengine as mge
import megengine.module as M

data = mge.tensor(np.array([-2,-1,0,1,2,]).astype(np.float32))
sigmoid = M.Sigmoid()
output = sigmoid(data)
with np.printoptions(precision=6):
    print(output.numpy())

输出：

[0.119203 0.268941 0.5      0.731059 0.880797]

forward(inputs)[源代码]¶

class megengine.module.activation.Softmax(axis=None)[源代码]¶

基类：megengine.module.module.Module

应用一个softmax函数。SoftMax定义为：

\[\text{Softmax}(x_{i}) = \frac{exp(x_i)}{\sum_j exp(x_j)}\]

应用softmax于一个n维输入张量，并重新放缩张量中的元素值，使得n维输出张量中所有元素的取值范围为 [0,1] 并且加和为1。

参数: axis – 沿着该轴应用softmax。默认情况下，softmax应用于最高维。

示例代码：

import numpy as np
import megengine as mge
import megengine.module as M

data = mge.tensor(np.array([-2,-1,0,1,2]).astype(np.float32))
softmax = M.Softmax()
output = softmax(data)
with np.printoptions(precision=6):
    print(output.numpy())

输出：

[0.011656 0.031685 0.086129 0.234122 0.636409]

forward(inputs)[源代码]¶

megengine.module.adaptive_pooling¶

class megengine.module.adaptive_pooling.AdaptiveAvgPool2d(oshp)[源代码]¶

基类：megengine.module.adaptive_pooling._AdaptivePoolNd

对输入数据进行2D平均池化。

例如，给定形状为 \((N, C, H, W)\) 的输入以及形为 \((kH, kW)\) 的 kernel_size ，该层产生形状为 \((N, C, H_{out}, W_{out})\) 的输出。生成过程描述如下：

\[out(N_i, C_j, h, w) = \frac{1}{kH * kW} \sum_{m=0}^{kH-1} \sum_{n=0}^{kW-1} input(N_i, C_j, stride[0] \times h + m, stride[1] \times w + n)\]

kernel_size 和 stride 可以通过输入和输出的形状推导出来：padding: (0, 0) stride: (floor(IH / OH), floor(IW / OW)) kernel_size: (IH - (OH - 1) * stride_h, IW - (OW - 1) * stride_w)

示例代码：

import numpy as np
import megengine as mge
import megengine.module as M

m = M.AdaptiveAvgPool2d((2, 2))
inp = mge.tensor(np.arange(0, 16).astype("float32").reshape(1, 1, 4, 4))
oup = m(inp)
print(oup.numpy())

输出：

[[[[ 2.5  4.5]
   [10.5 12.5]]]]

forward(inp)[源代码]¶

class megengine.module.adaptive_pooling.AdaptiveMaxPool2d(oshp)[源代码]¶

基类：megengine.module.adaptive_pooling._AdaptivePoolNd

对输入数据进行2D最大池化。

例如，给定形状为 \((N, C, H, W)\) 的输入以及形为 \((kH, kW)\) 的 kernel_size ，该层产生形状为 \((N, C, H_{out}, W_{out})\) 的输出。生成过程描述如下：

\[\begin{aligned} out(N_i, C_j, h, w) ={} & \max_{m=0, \ldots, kH-1} \max_{n=0, \ldots, kW-1} \text{input}(N_i, C_j, \text{stride[0]} \times h + m, \text{stride[1]} \times w + n) \end{aligned}\]

示例代码：

import numpy as np
import megengine as mge
import megengine.module as M

m = M.AdaptiveMaxPool2d((2, 2))
inp = mge.tensor(np.arange(0, 16).astype("float32").reshape(1, 1, 4, 4))
oup = m(inp)
print(oup.numpy())

输出：

[[[[ 5.  7.]
   [13. 15.]]]]

forward(inp)[源代码]¶

megengine.module.batch_matmul_activation¶

class megengine.module.batch_matmul_activation.BatchMatMulActivation(batch, in_features, out_features, bias=True, nonlinear_mode='IDENTITY', **kwargs)[源代码]¶

基类：megengine.module.module.Module

带激活的批矩阵乘（当前仅支持 relu 激活），输入均不支持转置。

forward(x)[源代码]¶

reset_parameters()[源代码]¶

返回类型: None

megengine.module.batchnorm¶

class megengine.module.batchnorm.BatchNorm1d(num_features, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True, freeze=False)[源代码]¶

基类：megengine.module.batchnorm._BatchNorm

在2D/3D张量上进行批标准化（batch normalization）。

参见 BatchNorm2d 以获取更多信息。

class megengine.module.batchnorm.BatchNorm2d(num_features, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True, freeze=False)[源代码]¶

基类：megengine.module.batchnorm._BatchNorm

在四维张量上进行批标准化（Batch Normalization）。

\[y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta\]

计算mini-batch各维的均值和标准差。其中， \(\gamma\) 和 \(\beta\) 是可学习的参数向量。

默认在训练过程中，该层持续估计其计算的均值和方差，随后在评估过程中用于标准化。运行中的估计保持默认的 momentum 值 0.9。

如果 track_running_stats 设置为 False ，此层将不会累计估计统计数据，而是把当前batch的统计数据应用在评估阶段。

注解

与在优化器类和传统动量概念中的 momentum 不同，这里统计特性的更新规则为： \(\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t\) ，其中 \(\hat{x}\) 是估计的统计特性， \(x_t\) 是新观测值。

因为批标准化是在 C 维进行的，(N, H, W) 切片上进行统计计算，所以通常将此方法称作空域批正则（Spatial Batch Normalization）。

参数

num_features (int) – 通常是形状为 \((N, C, H, W)\) 输入数据的 \(C\) 或者维度低于四维的输入的最高维。
eps (float) – 添加到分母的单个值，增加数值稳定性。默认：1e-5
momentum (float) – 用于计算 running_mean 和 running_var 的值。默认：0.9
affine (bool) – 单个布尔值，当设置为 True ，那么这个模块具有可学习的仿射（affine）参数。默认：True
track_running_stats (bool) – 当设置为 True，则这个模块跟踪运行时的不同batch的均值和方差。当设置为 False，该模块不跟踪这样的统计数据并在训练和eval模式下始终使用当前批统计数据。默认： True
freeze (bool) – 设置为True时，此模块不会更新运行运行时平均值和运行时方差，使用运行时平均值和方差而不是批次均值和批次方差来标准化输入。这个参数产生作用仅在使用 track_running_stats=True 初始化模块时有效并且模块处于训练模式。默认：False

示例代码：

import numpy as np
import megengine as mge
import megengine.module as M

# With Learnable Parameters
m = M.BatchNorm2d(4)
inp = mge.tensor(np.random.rand(1, 4, 3, 3).astype("float32"))
oup = m(inp)
print(m.weight.numpy().flatten(), m.bias.numpy().flatten())
# Without L`e`arnable Parameters
m = M.BatchNorm2d(4, affine=False)
oup = m(inp)
print(m.weight, m.bias)

输出：

[1. 1. 1. 1.] [0. 0. 0. 0.]
None None

class megengine.module.batchnorm.SyncBatchNorm(num_features, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True, freeze=False, group=<megengine.distributed.group.Group object>)[源代码]¶

基类：megengine.module.batchnorm._BatchNorm

使用同步批标准化。

forward(inp)[源代码]¶

megengine.module.concat¶

class megengine.module.concat.Concat[源代码]¶

基类：megengine.module.module.Module

一个与函数型concat功能相同的 Module 。可用 quantize_qat() 替换成 QATModule 型的 Concat 。

forward(inps, axis=0)[源代码]¶

megengine.module.conv¶

class megengine.module.conv.Conv1d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, conv_mode='CROSS_CORRELATION', compute_mode='DEFAULT')[源代码]¶

基类：megengine.module.conv._ConvNd

对输入张量进行一维卷积

例如，给定形状为 \((N, C_{\text{in}}, H)\) 的输入，该层生成形状为 \((N, C_{\text{out}}, H_{\text{out}})\) 的输出，生成过程如下：

\[\text{out}(N_i, C_{\text{out}_j}) = \text{bias}(C_{\text{out}_j}) + \sum_{k = 0}^{C_{\text{in}} - 1} \text{weight}(C_{\text{out}_j}, k) \star \text{input}(N_i, k)\]

其中 \(\star\) 是有效的1D互相关运算; \(N\) 是批大小; \(C\) 表示通道数; \(H\) 是以像素为单位输入平面的高度。

当 groups == in_channels 且 out_channels == K * in_channels ，其中 K 是正整数，该操作也被称为深度方向卷积（depthwise convolution）。

换句话说，对于形状为 \((N, C_{in}, H_{in})\) 的输入，可通过参数 \((in\_channels=C_{in}, out\_channels=C_{in} \times K, ..., groups=C_{in})\) ，构造一个深度方向乘法器 K。

参数

in_channels (int) – 输入数据中的通道数。
out_channels (int) – 输出数据中的通道数。
kernel_size (int) – 在空域维度权重的尺寸。如果 kernel_size 是一个 int ，则实际的卷积核尺寸是 (kernel_size, kernel_size) 。默认：1
stride (int) – 一维卷积运算的步长。默认：1
padding (int) – 输入数据空域维度两侧的填充（padding）大小。仅支持填充0值。默认：0
dilation (int) – 一维卷积运算的空洞（dilation）。默认：1
groups (int) – 将输入与输出通道分组的数目，用于进行分组卷积（grouped convolution）。当 groups 不为1时， in_channels 和 out_channels 必须可被 groups 整除，并且在权重形状的开始附加一个额外的维度。具体来说，权重的形状会变为 (groups, out_channel // groups, in_channels // groups, *kernel_size) 。默认：1
bias (bool) – 是否将偏置（bias）加入卷积的结果中。默认：True
conv_mode (str) – 支持 CROSS_CORRELATION。默认：CROSS_CORRELATION
compute_mode (str) – 若设定为 DEFAULT ，则对中间结果的精度没有特殊的要求。当设定为 FLOAT32 时，FLOAT32 将作为累加器和中间结果的精度，但只有当输入和输出的 dtype 是 float16 时才生效。

示例代码：

import numpy as np
import megengine as mge
import megengine.module as M

m = M.Conv1d(in_channels=3, out_channels=1, kernel_size=3)
inp = mge.tensor(np.arange(0, 24).astype("float32").reshape(2, 3, 4))
oup = m(inp)
print(oup.numpy().shape)

输出：

(2, 1, 2)

calc_conv(inp, weight, bias)[源代码]¶

forward(inp)[源代码]¶

class megengine.module.conv.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, conv_mode='CROSS_CORRELATION', compute_mode='DEFAULT')[源代码]¶

基类：megengine.module.conv._ConvNd

对输入张量进行二维卷积

例如，给定形状为 \((N, C_{\text{in}}, H, W)\) 的输入，该层生成形状为 \((N, C_{\text{out}}, H_{\text{out}}, W_{\text{out}})\) 的输出，生成过程如下：

\[\text{out}(N_i, C_{\text{out}_j}) = \text{bias}(C_{\text{out}_j}) + \sum_{k = 0}^{C_{\text{in}} - 1} \text{weight}(C_{\text{out}_j}, k) \star \text{input}(N_i, k)\]

其中 \(\star\) 是有效的2D互相关运算; \(N\) 是批大小; \(C\) 表示通道数; \(H\) 是以像素为单位输入平面的高度; \(W\) 是以像素为单位的平面宽度。

通常，输出的特征图的形状可以被下面的方式推导出来：

input: \((N, C_{\text{in}}, H_{\text{in}}, W_{\text{in}})\) output: \((N, C_{\text{out}}, H_{\text{out}}, W_{\text{out}})\) 其中

\[\text{H}_{out} = \lfloor \frac{\text{H}_{in} + 2 * \text{padding[0]} - \text{dilation[0]} * (\text{kernel_size[0]} - 1)}{\text{stride[0]}} + 1 \rfloor\]

\[\text{W}_{out} = \lfloor \frac{\text{W}_{in} + 2 * \text{padding[1]} - \text{dilation[1]} * (\text{kernel_size[1]} - 1)}{\text{stride[1]}} + 1 \rfloor\]

当 groups == in_channels 且 out_channels == K * in_channels ，其中 K 是正整数，该操作也被称为深度方向卷积（depthwise convolution）。

换句话说，对于形状为 \((N, C_{in}, H_{in}, W_{in})\) 的输入，可通过参数 \((in\_channels=C_{in}, out\_channels=C_{in} \times K, ..., groups=C_{in})\) ，构造一个深度方向乘法器 K。

参数

in_channels (int) – 输入数据中的通道数。
out_channels (int) – 输出数据中的通道数。
kernel_size (Union[int, Tuple[int, int]]) – 在空域维度权重的尺寸。如果 kernel_size 是一个 int ，则实际的卷积核尺寸是 (kernel_size, kernel_size) 。默认：1
stride (Union[int, Tuple[int, int]]) – 二维卷积运算的步长。默认：1
padding (Union[int, Tuple[int, int]]) – 输入数据空域维度两侧的填充（padding）大小。仅支持填充0值。默认：0
dilation (Union[int, Tuple[int, int]]) – 二维卷积运算的空洞（dilation）。默认：1
groups (int) – 将输入与输出通道分组的数目，用于进行分组卷积（grouped convolution）。当 groups 不为1时， in_channels 和 out_channels 必须可被 groups 整除，并且在权重形状的开始附加一个额外的维度。具体来说，权重的形状会变为 (groups, out_channel // groups, in_channels // groups, *kernel_size) 。默认：1
bias (bool) – 是否将偏置（bias）加入卷积的结果中。默认：True
conv_mode (str) – 支持 CROSS_CORRELATION。默认：CROSS_CORRELATION
compute_mode (str) – 若设定为 DEFAULT ，则对中间结果的精度没有特殊的要求。当设定为 FLOAT32 时，FLOAT32 将作为累加器和中间结果的精度，但只有当输入和输出的 dtype 是 float16 时才生效。

示例代码：

import numpy as np
import megengine as mge
import megengine.module as M

m = M.Conv2d(in_channels=3, out_channels=1, kernel_size=3)
inp = mge.tensor(np.arange(0, 96).astype("float32").reshape(2, 3, 4, 4))
oup = m(inp)
print(oup.numpy().shape)

输出：

(2, 1, 2, 2)

calc_conv(inp, weight, bias)[源代码]¶

forward(inp)[源代码]¶

class megengine.module.conv.ConvRelu2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, conv_mode='CROSS_CORRELATION', compute_mode='DEFAULT')[源代码]¶

基类：megengine.module.conv.Conv2d

融合了 Conv2d 和 relu 的 Module 。可以调用 quantize_qat() 替换成 QATModule 型的 ConvRelu2d .

forward(inp)[源代码]¶

class megengine.module.conv.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, conv_mode='CROSS_CORRELATION', compute_mode='DEFAULT')[源代码]¶

基类：megengine.module.conv._ConvNd

对输入张量进行二维转置卷积。

该模块也称为反卷积或微步卷积。 ConvTranspose2d 可以看作是 Conv2d 操作对自身输入的梯度。

卷积通常降低了输入的大小，而转置卷积工作方式完全相反，在相同的连接方式下，它对较小的输入进行放大来输出。

参数

in_channels (int) – 输入数据中的通道数。
out_channels (int) – 输出数据中的通道数。
kernel_size (Union[int, Tuple[int, int]]) – 在空域维度权重的尺寸。如果 kernel_size 是一个 int ，则实际的卷积核尺寸是 (kernel_size, kernel_size) 。默认：1
stride (Union[int, Tuple[int, int]]) – 二维卷积运算的步长。默认：1
padding (Union[int, Tuple[int, int]]) – 输入数据空域维度两侧的填充（padding）大小。仅支持填充0值。默认：0
dilation (Union[int, Tuple[int, int]]) – 二维卷积运算的空洞（dilation）。默认：1
groups (int) – 将输入与输出通道分组的数目，用于进行分组卷积（grouped convolution）。当 groups 不为1时， in_channels 和 out_channels 必须可被 groups 整除，并且在权重形状的开始附加一个额外的维度。具体来说，权重的形状会变为 (groups, out_channel // groups, in_channels // groups, *kernel_size)。默认：1
bias (bool) – 是否将偏置（bias）加入卷积的结果中。默认：True。
conv_mode (str) – 支持 CROSS_CORRELATION。默认：CROSS_CORRELATION
compute_mode (str) – 若设定为 DEFAULT ，则对中间结果的精度没有特殊的要求。当设定为 FLOAT32 时，FLOAT32 将作为累加器和中间结果的精度，但只有当输入和输出的 dtype 是 float16 时才生效。

forward(inp)[源代码]¶

class megengine.module.conv.LocalConv2d(in_channels, out_channels, input_height, input_width, kernel_size, stride=1, padding=0, dilation=1, groups=1, conv_mode='CROSS_CORRELATION')[源代码]¶

基类：megengine.module.conv.Conv2d

在输入的4D张量上使用多个untied kernel进行空域卷积。它也被称为局部连接层。

参数

in_channels (int) – 输入数据中的通道数。
out_channels (int) – 输出数据中的通道数。
input_height (int) – 输入图像的高度。
input_width (int) – 输入图像的宽度。
kernel_size (Union[int, Tuple[int, int]]) – 在空域维度权重的尺寸。如果 kernel_size 是一个 int ，则实际的卷积核尺寸是 (kernel_size, kernel_size) 。默认：1
stride (Union[int, Tuple[int, int]]) – 二维卷积运算的步长。默认：1
padding (Union[int, Tuple[int, int]]) – 输入数据空域维度两侧的填充（padding）大小。仅支持填充0值。默认：0
groups (int) – 将输入与输出通道分组的数目，用于进行分组卷积（grouped convolution）。当 groups 不为1时， in_channels 和 out_channels 必须可被 groups 整除。权重的形状会变为 (groups, out_channel // groups, in_channels // groups, *kernel_size)。

forward(inp)[源代码]¶

megengine.module.conv_bn¶

class megengine.module.conv_bn.ConvBn2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, conv_mode='CROSS_CORRELATION', compute_mode='DEFAULT', eps=1e-05, momentum=0.9, affine=True, track_running_stats=True)[源代码]¶

基类：megengine.module.conv_bn._ConvBnActivation2d

融合了 Conv2d 和 BatchNorm2d 的 Module 。可以调用 quantize_qat() 替换成 QATModule 型的 ConvBn2d .

forward(inp)[源代码]¶

class megengine.module.conv_bn.ConvBnRelu2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, conv_mode='CROSS_CORRELATION', compute_mode='DEFAULT', eps=1e-05, momentum=0.9, affine=True, track_running_stats=True)[源代码]¶

基类：megengine.module.conv_bn._ConvBnActivation2d

融合了 Conv2d ，BatchNorm2d 和 relu 的 Module 。可以调用 quantize_qat() 替换成 QATModule 型的 ConvBnRelu2d .

forward(inp)[源代码]¶

megengine.module.dropout¶

class megengine.module.dropout.Dropout(drop_prob=0.0)[源代码]¶

基类：megengine.module.module.Module

训练时，依概率 \(drop\_prob\) 随机地将某些输入元素置为0。在大型网络中有广泛应用，可以有效防止过拟合。注意，我们只在训练过程中执行dropout操作。输出张量也会通过因子 \(\frac{1}{1 - p}\) 进行放缩（即相乘）。在预测（inference）阶段， Dropout 与 Identity 等效。

参数: drop_prob – 每个元素被置为0的概率

forward(inputs)[源代码]¶

megengine.module.elemwise¶

class megengine.module.elemwise.Elemwise(method)[源代码]¶

基类：megengine.module.module.Module

相当于elemwise算子的 Module 。可调用 quantize_qat() 替换成 QATModule 型的 class:~.qat.elemwise.Elemwise 。

参数

method –

elemwise方法，会进行标准elemwise浮点数运算。支持如下字符串：

”ADD”: a + b
”FUSE_ADD_RELU”: max(x+y, 0)
”MUL”: x * y
”MIN”: min(x, y)
”MAX”: max(x, y)
”SUB”: x - y
”TRUE_DIV”: x / y
”FUSE_ADD_SIGMOID”: sigmoid(x + y)
”FUSE_ADD_TANH”: tanh(x + y)
”RELU”: x > 0 ? x : 0
”ABS”: x > 0 ? x : -x
”SIGMOID”: sigmoid(x)
”EXP”: exp(x)
”TANH”: tanh(x)
”FUSE_MUL_ADD3”: x * y + z
”FAST_TANH”: x * (27. + x * x) / (27. + 9. * x * x)
”NEGATE”: -x
”ACOS”: acos(x)
”ASIN”: asin(x)
”CEIL”: ceil(x)
”COS”: cos(x)
”EXPM1”: expm1(x)
”FLOOR”: floor(x)
”LOG”: log(x)
”LOG1P”: log1p(x)
”SIN”: sin(x)
”ROUND”: round(x)
”ERF”: erf(x)
”ERFINV”: erfinv(x)
”ERFC”: erfc(x)
”ERFCINV”: erfcinv(x)
”ABS_GRAD”: abs_grad
”FLOOR_DIV”: floor_div
”MOD”: mod
”SIGMOID_GRAD”: sigmoid_grad
”SWITCH_GT0”: switch_gt0
”TANH_GRAD”: tanh_grad
”LEQ”: less
”LEQ”: leq
”EQ”: equal
”POW”: pow
”LOG_SUM_EXP”: log_sum_exp
”FAST_TANH_GRAD”: fast_tanh_grad
”ATAN2”: atan2
”COND_LEQ_MOV”: cond_leq_mov
”H_SWISH”: h_swish
”FUSE_ADD_H_SWISH”: h_swish(x+y)
”H_SWISH_GRAD”: h_swish_grad
”AND”: bool binary: x && y
”OR”: bool binary: x || y
”XOR”: bool binary: x ^ y
”NOT”: bool unary: ~x

forward(*inps)[源代码]¶

megengine.module.embedding¶

class megengine.module.embedding.Embedding(num_embeddings, embedding_dim, padding_idx=None, max_norm=None, norm_type=None, initial_weight=None, freeze=False)[源代码]¶

基类：megengine.module.module.Module

一个简单的查询表，存储具有固定大小的词向量（embedding）于固定的词典中。

该模块通常用于存储词向量（word embeddings），并使用索引来检索。输入索引列表到模块中，则输出对应的词向量。索引值应小于num_embeddings。

参数

num_embeddings (int) – 词向量字典的大小。
embedding_dim (int) – 每个词向量的大小。
padding_idx (Optional[int]) – 应设置为None，目前不支持。
max_norm (Optional[float]) – 应设置为None，目前不支持。
norm_type (Optional[float]) – 应设置为None，目前不支持。
initial_weight (Optional[Parameter]) – 该模块的可学习权重，形状为(num_embeddings, embedding_dim) 。

示例代码：

import numpy as np
import megengine as mge
import megengine.module as M
weight = mge.tensor(np.array([(1.2,2.3,3.4,4.5,5.6)], dtype=np.float32))
data = mge.tensor(np.array([(0,0)], dtype=np.int32))

embedding = M.Embedding(1, 5, initial_weight=weight)
output = embedding(data)
with np.printoptions(precision=6):
    print(output.numpy())

输出：

[[[1.2 2.3 3.4 4.5 5.6]
  [1.2 2.3 3.4 4.5 5.6]]]

forward(inputs)[源代码]¶

classmethod from_pretrained(embeddings, freeze=True, padding_idx=None, max_norm=None, norm_type=None)[源代码]¶

从给定的2维FloatTensor创建词向量实例。

参数

embeddings (Parameter) – 包含词向量权重的张量。
freeze (Optional[bool]) – 如果为 True ，则在学习过程中不更新权重。默认：True
padding_idx (Optional[int]) – 应设置为None，目前不支持。
max_norm (Optional[float]) – 应设置为None，目前不支持。
norm_type (Optional[float]) – 应设置为None，目前不支持。

示例代码：

import numpy as np
import megengine as mge
import megengine.module as M
weight = mge.tensor(np.array([(1.2,2.3,3.4,4.5,5.6)], dtype=np.float32))
data = mge.tensor(np.array([(0,0)], dtype=np.int32))

embedding = M.Embedding.from_pretrained(weight, freeze=False)
output = embedding(data)
print(output.numpy())

输出：

[[[1.2 2.3 3.4 4.5 5.6]
  [1.2 2.3 3.4 4.5 5.6]]]

reset_parameters()[源代码]¶

返回类型: None

megengine.module.identity¶

class megengine.module.identity.Identity[源代码]¶

基类：megengine.module.module.Module

恒等占位符，输出和输入相等

forward(x)[源代码]¶

megengine.module.init¶

megengine.module.init.calculate_correct_fan(tensor, mode)[源代码]¶

计算对于给定的权重张量fan_in或fan_out值，这取决于给定的 mode。

请参阅 calculate_fan_in_and_fan_out() 了解详情。

参数

tensor (Tensor) – NCHW 格式的权重张量。
mode (str) – 'fan_in' 或 'fan_out'。

返回类型

float

megengine.module.init.calculate_fan_in_and_fan_out(tensor)[源代码]¶

对于给定的权重张量，计算fan_in / fan_out值。此函数假定输入张量存储格式为NCHW。

参数: tensor (Tensor) – NCHW 格式的权重张量。
返回类型: Tuple[float, float]

megengine.module.init.calculate_gain(nonlinearity, param=None)[源代码]¶

对于给定的非线性函数返回一个推荐的增益值（见下表）。

非线性	增益
Linear / Identity	\(1\)
Conv{1,2,3}D	\(1\)
Sigmoid	\(1\)
Tanh	\(\frac{5}{3}\)
ReLU	\(\sqrt{2}\)
Leaky Relu	\(\sqrt{\frac{2}{1 + {\text{negative}_\text{slope}}^2}}\)

参数

nonlinearity (str) – 非线性函数的名称。
param (Union[int, float, None]) – leaky_relu的可选参数。只有当 nonlinearity 是 “leaky_relu” 时有效。

返回类型

float

megengine.module.init.fill_(tensor, val)[源代码]¶

将 val 填入到给定的 tensor。

参数

tensor (Tensor) – 待初始化的n维张量。
val (Union[float, int]) – 在整个张量上填入的值。

返回类型

None

megengine.module.init.msra_normal_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu')[源代码]¶

从 \(\mathcal{N}(0, \text{std}^2)\) 中随机采样得到值，填入 tensor ，式中，

\[\text{std} = \sqrt{\frac{2}{(1 + a^2) \times \text{fan_in}}}\]

参考 “Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification” 获取详细说明

参数

tensor (Tensor) – 待初始化的n维张量
a (float) – 用于计算leaky_relu增益的可选参数。请参阅： calculate_gain() 了解详情。
mode (str) – 'fan_in' 或 'fan_out' ，用于计算 \(gain\) ，是 \(gain\) 的比例因子。请参阅 calculate_fan_in_and_fan_out() 了解详情。
nonlinearity (str) – 用于计算 \(gain\) 的非线性函数的名称。请参阅 calculate_gain() 了解详情。

返回类型

None

megengine.module.init.msra_uniform_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu')[源代码]¶

使用 \(\mathcal{U}(-\text{bound}, \text{bound})\) 的随机采样值填入 tensor

\[\text{bound} = \sqrt{\frac{6}{(1 + a^2) \times \text{fan_in}}}\]

参考 “Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification” 获取详细说明

参数

tensor (Tensor) – 待初始化的n维张量。
a (float) – 用于计算leaky_relu增益的可选参数。请参阅： calculate_gain() 了解详情。
mode (str) – `` ‘fan_in’`` 或 ' fan_out' ，用于计算 \(gain\) ，为 \(bound\) 的比例因子。请参阅： calculate_fan_in_and_fan_out() 了解详情。
nonlinearity (str) – 用于计算 \(gain\) 的非线性函数的名称。请参阅 calculate_gain() 了解详情。

返回类型

None

megengine.module.init.normal_(tensor, mean=0.0, std=1.0)[源代码]¶

使用从正态分布 \(\mathcal{N}(\text{mean}, \text{std}^2)\) 中的随机采样值填入给定的 tensor

参数

tensor (Tensor) – 待初始化的n维张量。
mean (float) – 正态分布的均值。
std (float) – 正态分布的标准差。

返回类型

None

megengine.module.init.ones_(tensor)[源代码]¶

使用标量值 1 填入给定的 tensor。

参数: tensor (Tensor) – 待初始化的n维张量。
返回类型: None

megengine.module.init.uniform_(tensor, a=0.0, b=1.0)[源代码]¶

使用从均匀分布 \(\mathcal{U}(\text{a}, \text{b})\) 中的随机采样值填入给定的 tensor。

参数

tensor (Tensor) – 待初始化的n维张量。
a (float) – 采样区间的下界。
b (float) – 采样区间的上界。

返回类型

None

megengine.module.init.xavier_normal_(tensor, gain=1.0)[源代码]¶

使用从 \(\mathcal{N}(0, \text{std}^2)\) 中随机采样值填入 tensor 。其中，

\[\text{std} = \text{gain} \times \sqrt{\frac{2}{\text{fan_in} + \text{fan_out}}}\]

又称Glorot初始化。详细的说明可以参考 “Understanding the difficulty of training deep feedforward neural networks” 。

参数

tensor (Tensor) – 待初始化的n维张量。
gain (float) – \(std\) 的比例因子。

返回类型

None

megengine.module.init.xavier_uniform_(tensor, gain=1.0)[源代码]¶

使用从 \(\mathcal{U}(-a, a)\) 中随机采样值填入 tensor 。其中，

\[a = \text{gain} \times \sqrt{\frac{6}{\text{fan_in} + \text{fan_out}}}\]

又称Glorot初始化。详细的说明可以参考 “Understanding the difficulty of training deep feedforward neural networks” 。

参数

tensor (Tensor) – 待初始化的n维张量。
gain (float) – \(a\) 的比例因子。

返回类型

None

megengine.module.init.zeros_(tensor)[源代码]¶

使用标量值 0 填入给定 tensor。

参数: tensor (Tensor) – 待初始化的n维张量。
返回类型: None

megengine.module.linear¶

class megengine.module.linear.Linear(in_features, out_features, bias=True, **kwargs)[源代码]¶

基类：megengine.module.module.Module

对输入进行线性变换。例如，若有输入x，则输出y为：

\[y = xW^T + b\]

其中 \(y_i= \sum_j W_{ij} x_j + b_i\)

参数

in_features (int) – 各输入样本的大小。
out_features (int) – 各输出样本的大小。
bias (bool) – 如果设置为 False ，则该层不会学习加性偏置。默认： True

示例代码：

import numpy as np
import megengine as mge
import megengine.module as M

m = M.Linear(in_features=3, out_features=1)
inp = mge.tensor(np.arange(0, 6).astype("float32").reshape(2, 3))
oup = m(inp)
print(oup.numpy().shape)

输出：

(2, 1)

forward(x)[源代码]¶

reset_parameters()[源代码]¶

返回类型: None

megengine.module.module¶

class megengine.module.module.Module[源代码]¶

基类：object

Module基类。

apply(fn)[源代码]¶

对当前模块中的所有模块应用函数 fn，包括当前模块本身。

参数: fn (Callable[[Module], Any]) – 多个模块上要应用的函数。
返回类型: None

buffers(recursive=True, **kwargs)[源代码]¶

返回该模块中对于buffers的一个可迭代对象。

Buffer被定义为是 Tensor 且不是 Parameter

参数: recursive (bool) – 如果为 True ，则返回所有当前模块中的buffer，否则只返回属于该模块的直接属性。
返回类型: Iterable[Tensor]

children(**kwargs)[源代码]¶

返回一个可迭代对象，可遍历所有属于当前模块的直接属性的子模块。

返回类型: Iterable[Module]

disable_quantize(value=True)[源代码]¶: 设置 module 的 quantize_diabled 属性，并返回 module 。可以作为装饰器使用。

eval()[源代码]¶

当前模块中所有模块的 training 属性（包括自身）置为 False ，并将其切换为推理模式。请参阅 train() 了解详情。

返回类型: None

abstract forward(inputs)[源代码]¶

load_state_dict(state_dict, strict=True)[源代码]¶

向当前模块中加载由 state_dict() 创建的给定字典。若 strict 为 True ， state_dict() 的键则必须与 state_dict() 返回的键准确匹配。

为了处理复杂情况，用户可以传入闭包 Function[key: str, var: Tensor] -> Optional[np.ndarray] 作为 state_dict 。例如，欲加载除了最后线性分类器外的所有部分：

state_dict = {...}  #  Dict[str, np.ndarray]
model.load_state_dict({
    k: None if k.startswith('fc') else v
    for k, v in state_dict.items()
}, strict=False)

这里返回 None 意味着忽略参数 k 。

为了防止形状不匹配（例如加载PyTorch权重），我们可以在加载之前重塑（reshape）：

state_dict = {...}
def reshape_accordingly(k, v):
    return state_dict[k].reshape(v.shape)
model.load_state_dict(reshape_accordingly)

我们还可以进行原位重初始化或修剪（pruning）：

def reinit_and_pruning(k, v):
    if 'bias' in k:
        M.init.zero_(v)
    if 'conv' in k:
        return v.numpy() * (np.abs(v.numpy()) > 1e-3).astype("float32)
model.load_state_dict(reinit_and_pruning, strict=False)

modules(**kwargs)[源代码]¶

返回一个可迭代对象，可以遍历当前模块中的所有模块，包括其本身。

返回类型: Iterable[Module]

named_buffers(prefix=None, recursive=True, **kwargs)[源代码]¶

返回可遍历模块中 key 与 buffer 的键值对的可迭代对象，其中 key 为从该模块至 buffer 的点路径（dotted path）。

Buffer被定义为是 Tensor 且不是 Parameter

参数

prefix (Optional[str]) – 加在每个键（key）前的前缀。
recursive (bool) – 如果为 True ，则返回所有当前模块中的buffer，否则只返回属于该模块的直接属性。

返回类型

Iterable[Tuple[str, Buffer]]

named_children(**kwargs)[源代码]¶

返回可迭代对象，可以遍历属于当前模块的直接属性的所有子模块（submodule）与键（key）组成的”key-submodule”对，其中’key’是子模块对应的属性名。

返回类型: Iterable[Tuple[str, Module]]

named_modules(prefix=None, **kwargs)[源代码]¶

返回可迭代对象，可以遍历当前模块包括自身在内的所有其内部模块所组成的key-module键-模块对，其中’key’是从当前模块到各子模块的点路径（dotted path）。

参数: prefix (Optional[str]) – 加在路径前的前缀。
返回类型: Iterable[Tuple[str, Module]]

named_parameters(prefix=None, recursive=True, **kwargs)[源代码]¶

返回一个可迭代对象，可以遍历当前模块中key与 Parameter 组成的键值对。其中 key 是从模块到 Parameter 的点路径（dotted path）。

参数

prefix (Optional[str]) – 加在每个键（key）前的前缀。
recursive (bool) – 如果为 True ，则返回在此模块内的所有 Parameter ; 否则，只返回属于当前模块直接属性的 Parameter 。

返回类型

Iterable[Tuple[str, Parameter]]

parameters(recursive=True, **kwargs)[源代码]¶

返回一个可迭代对象，遍历当前模块中的所有 Parameter

参数: recursive (bool) – 如果为 True ，则返回在此模块内的所有 Parameter ; 否则，只返回属于当前模块直接属性的 Parameter 。
返回类型: Iterable[Parameter]

register_forward_hook(hook)[源代码]¶

注册一个回调函数以处理转发结果。hook 应该是一个函数，它可以接受 module， inputs 和 outputs 作为输入，然后返回修改后的`outputs` 或者 None。

此方法返回具有 〜.HookHandler.remove() 接口的句柄来删除回调函数。

返回类型: HookHandler

register_forward_pre_hook(hook)[源代码]¶

注册一个回调函数处理前向输入。 hook 应该是一个函数。

参数: hook (Callable) – 一个接受 module and inputs 作为输入的函数，然后返回

一个被更改的 inputs 或者 None。:rtype: HookHandler :return: 一个具有 remove() 接口的句柄来删除回调函数。

replace_param(params, start_pos, seen=None)[源代码]¶: 用 param 替换模块的参数，被 ParamPack 用来加速多机训练。

1.0 版后已移除.

state_dict(rst=None, prefix='', keep_var=False)[源代码]¶

train(mode=True, recursive=True)[源代码]¶

将当前模块，及其内部所有模块的训练模式置为 mode 。该方法将这些模块的 training 属性便捷地置为 mode 。但只对特定种类的模块起作用（比如： BatchNorm2d ， Dropout ， Observer ）。

参数

mode (bool) – 要在多个模块中设置的训练模式。
recursive (bool) – 是否递归调用子模块的 train() 。

返回类型

None

zero_grad()[源代码]¶

将所有参数的梯度置0。

1.0 版后已移除.

返回类型: None

megengine.module.normalization¶

class megengine.module.normalization.GroupNorm(num_groups, num_channels, eps=1e-05, affine=True)[源代码]¶

基类：megengine.module.module.Module

GroupNorm 的简单实现。当前仅支持四维张量。参考文献：https://arxiv.org/pdf/1803.08494.pdf。

forward(x)[源代码]¶

reset_parameters()[源代码]¶

class megengine.module.normalization.InstanceNorm(num_channels, eps=1e-05, affine=True)[源代码]¶

基类：megengine.module.module.Module

InstanceNorm 的简单实现。当前仅支持四维张量。参考文献：https://arxiv.org/abs/1607.08022。请注意当 num_groups=num_channels 时 InstanceNorm 与 GroupNorm 等价。

forward(x)[源代码]¶

reset_parameters()[源代码]¶

class megengine.module.normalization.LayerNorm(num_channels, eps=1e-05, affine=True)[源代码]¶

基类：megengine.module.module.Module

LayerNorm 的简单实现。当前仅支持四维张量。参考文献：https://arxiv.org/pdf/1803.08494.pdf。请注意当 num_groups=1 时 LayerNorm 与 GroupNorm 等价。

forward(x)[源代码]¶

reset_parameters()[源代码]¶

megengine.module.pooling¶

class megengine.module.pooling.AvgPool2d(kernel_size, stride=None, padding=0)[源代码]¶

基类：megengine.module.pooling._PoolNd

对输入数据进行2D平均池化。

例如，给定形状为 \((N, C, H, W)\) 的输入以及形为 \((kH, kW)\) 的 kernel_size ，该层产生形状为 \((N, C, H_{out}, W_{out})\) 的输出。生成过程描述如下：

\[out(N_i, C_j, h, w) = \frac{1}{kH * kW} \sum_{m=0}^{kH-1} \sum_{n=0}^{kW-1} input(N_i, C_j, stride[0] \times h + m, stride[1] \times w + n)\]

若 padding 非零，则输入数据会被隐式地在两边用零值进行填充（pad) padding 个点。

参数

kernel_size (Union[int, Tuple[int, int]]) – 窗的大小。
stride (Union[int, Tuple[int, int], None]) – 窗的步长。默认值是 kernel_size。
padding (Union[int, Tuple[int, int]]) – 对两边进行隐式的零值填充尺寸。

示例代码：

import numpy as np
import megengine as mge
import megengine.module as M

m = M.AvgPool2d(kernel_size=3, stride=1, padding=0)
inp = mge.tensor(np.arange(0, 16).astype("float32").reshape(1, 1, 4, 4))
oup = m(inp)
print(oup.numpy())

输出：

[[[[ 5.  6.]
   [ 9. 10.]]]]

forward(inp)[源代码]¶

class megengine.module.pooling.MaxPool2d(kernel_size, stride=None, padding=0)[源代码]¶

基类：megengine.module.pooling._PoolNd

对输入数据进行2D最大值池化（max pooling）。

例如，给定形状为 \((N, C, H, W)\) 的输入以及形为 \((kH, kW)\) 的 kernel_size ，该层产生形状为 \((N, C, H_{out}, W_{out})\) 的输出。生成过程描述如下：

\[\begin{aligned} out(N_i, C_j, h, w) ={} & \max_{m=0, \ldots, kH-1} \max_{n=0, \ldots, kW-1} \text{input}(N_i, C_j, \text{stride[0]} \times h + m, \text{stride[1]} \times w + n) \end{aligned}\]

若 padding 非零，则输入数据会被隐式地在两边用零值进行填充（pad) padding 个点。

参数

kernel_size (Union[int, Tuple[int, int]]) – 从中取最大值的窗口大小。
stride (Union[int, Tuple[int, int], None]) – 窗的步长。默认值是 kernel_size。
padding (Union[int, Tuple[int, int]]) – 对两边进行隐式的零值填充尺寸。

示例代码：

import numpy as np
import megengine as mge
import megengine.module as M

m = M.MaxPool2d(kernel_size=3, stride=1, padding=0)
inp = mge.tensor(np.arange(0, 16).astype("float32").reshape(1, 1, 4, 4))
oup = m(inp)
print(oup.numpy())

输出：

[[[[10. 11.]
   [14. 15.]]]]

forward(inp)[源代码]¶

megengine.module.quant_dequant¶

class megengine.module.quant_dequant.DequantStub[源代码]¶

基类：megengine.module.module.Module

仅返回输入的辅助 Module 。可调用 quantize_qat() 替换成 QATModule 型的 DequantStub 。

forward(inp)[源代码]¶

class megengine.module.quant_dequant.QuantStub[源代码]¶

基类：megengine.module.module.Module

仅返回输入的辅助 Module 。可调用 quantize_qat() 替换成 QATModule 型的 QuantStub 。

forward(inp)[源代码]¶

megengine.module.sequential¶

class megengine.module.sequential.Sequential(*args)[源代码]¶

基类：megengine.module.module.Module

一个序列容器。多个模块会按在构造函数中传入的顺序加到该容器中。或者，也可以使用模块组成的有序字典来存入容器。

为了便于理解，这里提供一个小例子：

示例代码：

import numpy as np
import megengine as mge
import megengine.module as M
import megengine.functional as F
from collections import OrderedDict

batch_size = 64
data = mge.tensor(np.zeros((batch_size, 28 * 28)), dtype=np.float32)
label = mge.tensor(np.zeros(batch_size,), dtype=np.int32)

net0 = M.Sequential(
        M.Linear(28 * 28, 320),
        M.Linear(320, 10)
    )
pred0 = net0(data)

modules = OrderedDict()
modules["fc0"] = M.Linear(28 * 28, 320)
modules["fc1"] = M.Linear(320, 10)
net1 = M.Sequential(modules)
pred1 = net1(data)

forward(inp)[源代码]¶

property layer_values¶