megengine.module

Float Module

Containers

Module

Module 基类。

Sequential

A sequential container.

Initialization

fill_

val 填入到给定的 tensor

zeros_

使用标量值 0 填入给定 tensor

ones_

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

uniform_

Fills the given tensor with random value sampled from uniform distribution \(\mathcal{U}(\text{a}, \text{b})\).

normal_

Fills the given tensor with random value sampled from normal distribution \(\mathcal{N}(\text{mean}, \text{std}^2)\).

calculate_gain

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

calculate_fan_in_and_fan_out

Calculates fan_in / fan_out value for given weight tensor.

calculate_correct_fan

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

xavier_uniform_

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

xavier_normal_

Fills tensor with random values sampled from \(\mathcal{N}(0, \text{std}^2)\) where

msra_uniform_

Fills tensor wilth random values sampled from \(\mathcal{U}(-\text{bound}, \text{bound})\) where

msra_normal_

Fills tensor wilth random values sampled from \(\mathcal{N}(0, \text{std}^2)\) where

General operations

Elemwise

进行 elemwise 运算的 Module.

Concat

进行 concat 运算的 Module.

Convolution Layers

Conv1d

对输入张量进行一维卷积

Conv2d

对输入张量进行二维卷积

Conv3d

对输入 tensor 进行三维卷积

ConvTranspose2d

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

ConvTranspose3d

Applies a 3D transposed convolution over an input tensor.

LocalConv2d

Applies a spatial convolution with untied kernels over an groupped channeled input 4D tensor.

DeformableConv2d

可变形卷积

SlidingWindow

Apply a sliding window to input tensor and copy content in the window to corresponding output location.

Pooling layers

AvgPool2d

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

MaxPool2d

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

AdaptiveAvgPool2d

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

AdaptiveMaxPool2d

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

DeformablePSROIPooling

Non-linear Activations

Sigmoid

对每个元素应用函数:

Softmax

Applies a softmax function.

ReLU

对每个元素应用函数:

LeakyReLU

对每个元素应用函数:

PReLU

对每个元素应用函数:

Normalization Layers

BatchNorm1d

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

BatchNorm2d

在四维张量上进行批标准化(Batch Normalization)。

SyncBatchNorm

对于分布式训练执行组同步Batch Normalization操作。

GroupNorm

对于GroupNorm的简单实现。将Channel方向分Group,然后每个Group内做归一化,算(C//G) Hw的均值;这样与Batchsize无关,不受其约束

InstanceNorm

InstanceNorm的简单实现。对Channel内做归一化,算H*W的均值,用在风格化迁移;因为在图像风格化中,生成结果主要依赖于某个图像实例,所以对整个Batch归一化不适合图像风格化中,因而对HW做归一化。可以加速模型收敛,并且保持每个图像实例之间的独立。

LayerNorm

层归一化的简单实现,对于Channel方向做归一化,算CHW的均值,主要对RNN作用明显。

Linear Layers

Identity

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

Linear

Applies a linear transformation to the input.

Dropout Layers

Dropout

Randomly sets input elements to zeros with the probability \(drop\_prob\) during training.

Sparse Layers

Embedding

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

Fused operations

ConvRelu2d

融合 Conv2drelu 运算的 Module.

ConvBn2d

融合 Conv2dBatchNorm2d 运算的 Module.

ConvBnRelu2d

融合 Conv2d, BatchNorm2drelu 运算的 Module.

BatchMatMulActivation

带激活的批矩阵乘 matmul (当前仅支持 relu 激活),输入均不支持转置。

SlidingWindow

Apply a sliding window to input tensor and copy content in the window to corresponding output location.

SlidingWindowTranspose

Opposite opration of SlidingWindow, sum over the sliding windows on the corresponding input location.

Quantization

QuantStub

仅负责返回输入的辅助 Module.

DequantStub

仅负责返回输入的辅助 Module.

QAT Module

Containers

QATModule

浮点数量化相关 Module 的基类。主要用于 QAT 和 Calibration.

Operations

Linear

A QATModule version of Linear.

Elemwise

支持 QAT 的 QATModule 用来执行 elemwise 计算.

Concat

支持 QAT 的 QATModule 用来执行 concat 计算.

Conv2d

支持 QAT 的 QATModule Conv2d.

ConvRelu2d

支持 QAT 的 QATModule 包括 Conv2drelu.

ConvBn2d

支持 QAT 的 QATModule 融合了 Conv2dBatchNorm2d.

ConvBnRelu2d

支持 QAT 的 QATModule 融合了 Conv2d, BatchNorm2drelu.

BatchMatMulActivation

支持 QAT 的 QATModule BatchMatMulActivation.

QuantStub

仅返回输入的辅助 QATModule, 但在转换为 QuantizedModule 之后,会对输入进行量化。

DequantStub

仅返回输入的辅助 QATModule, 但在转换为 QuantizedModule 之后,会对输入进行反量化。

Quantized Module

QuantizedModule

量化版本 Module 的基类。应从 QATModule 转换而来,不支持训练。

Operations

Linear

量化(quantized)版本 Linear

Elemwise

量化(quantized)版本 Elemwise

Concat

量化版本 concatQuantizedModule, 仅用于推理阶段。

Conv2d

量化(quantized)版本 Conv2d

ConvRelu2d

量化(quantized)版本 ConvRelu2d.

ConvBn2d

量化(quantized)版本 ConvBn2d

ConvBnRelu2d

量化(quantized)版本 ConvBnRelu2d.

BatchMatMulActivation

BatchMatMulActivation 的量化版本。

QuantStub

量化(quantized)版本 QuantStub,可将输入转化为量化数据类型。

DequantStub

量化(quantized)版本 DequantStub,可将量化后的输入重置为float32类型。