megengine.module package

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}\]

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.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.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_meanrunning_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(), m.bias.numpy())
# 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.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\) 是以像素为单位输入平面的高度; :math:`W`是以像素为单位的平面宽度。

groups == in_channelsout_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_channelsout_channels 必须可被 groups 整除,并且在权重形状的开始附加一个额外的维度。具体来说,权重的形状会变为 (groups, out_channel // groups, in_channels // groups, *kernel_size) 。 默认:1

  • bias (bool) – 是否将偏置(bias)加入卷积的结果中。默认:True

  • conv_mode (str) – 支持 CROSS_CORRELATION 或者 CONVOLUTION。默认: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.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_channelsout_channels 必须可被 groups 整除,并且在权重形状的开始附加一个额外的维度。具体来说,权重的形状会变为 (groups, out_channel // groups, in_channels // groups, *kernel_size)。 默认:1

  • bias (bool) – 是否将偏置(bias)加入卷积的结果中。默认:True。

  • conv_mode (str) – 支持 CROSS_CORRELATION 或者 CONVOLUTION。默认: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)阶段, DropoutIdentity 等效。

参数

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),(0.1,1.1,2.1,3.1,4.1)], dtype=np.float32))
data = mge.tensor(np.array([(0,1,1),(1,0,1),(0,0,1)], dtype=np.int32))

embedding = M.Embedding(2, 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]
  [0.1 1.1 2.1 3.1 4.1]
  [0.1 1.1 2.1 3.1 4.1]]

 [[0.1 1.1 2.1 3.1 4.1]
  [1.2 2.3 3.4 4.5 5.6]
  [0.1 1.1 2.1 3.1 4.1]]

 [[1.2 2.3 3.4 4.5 5.6]
  [1.2 2.3 3.4 4.5 5.6]
  [0.1 1.1 2.1 3.1 4.1]]]
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),(0.1,1.1,2.1,3.1,4.1)], dtype=np.float32))
data = mge.tensor(np.array([(0,1,1),(1,0,1),(0,0,1)], 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]
  [0.1 1.1 2.1 3.1 4.1]
  [0.1 1.1 2.1 3.1 4.1]]

 [[0.1 1.1 2.1 3.1 4.1]
  [1.2 2.3 3.4 4.5 5.6]
  [0.1 1.1 2.1 3.1 4.1]]

 [[1.2 2.3 3.4 4.5 5.6]
  [1.2 2.3 3.4 4.5 5.6]
  [0.1 1.1 2.1 3.1 4.1]]]
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.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)[源代码]

设置 modulequantize_diabled 属性,并返回 module 。可以作为装饰器使用。

eval()[源代码]

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

返回类型

None

abstract forward(inputs)[源代码]
load_state_dict(state_dict, strict=True)[源代码]

向当前模块中加载由 state_dict() 创建的给定字典。若 strictTruestate_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 应该是一个函数,它可以接受 moduleinputsoutputs 作为输入,然后返回修改后的`outputs` 或者 None

此方法返回具有:meth:〜.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 。但只对特定种类的模块起作用(比如: BatchNorm2dDropoutObserver )。

参数
  • mode (bool) – 要在多个模块中设置的训练模式。

  • recursive (bool) – 是否递归调用子模块的 train()

返回类型

None

zero_grad()[源代码]

将所有参数的梯度置0。

1.0 版后已移除.

返回类型

None

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, 1, 28, 28)), dtype=np.float32)
label = mge.tensor(np.zeros(batch_size,), dtype=np.int32)

data = data.reshape(batch_size, -1)
net0 = M.Sequential(
        M.Linear(28 * 28, 320),
        M.Linear(320, 10)
    )
pred0 = net0(data)

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

pred1 = net1(data)
forward(inp)[源代码]
property layer_values