megengine.module#

>>> import megengine.module as M

参见

浮点模块#

容器#

Module

Module 基类。

Sequential

一种序列化容器。

一般操作#

Elemwise

进行 elemwise 运算的 Module.

Concat

进行 concat 运算的 Module.

卷积层#

Conv1d

对输入张量进行一维卷积

Conv2d

对输入张量进行二维卷积

Conv3d

对输入 tensor 进行三维卷积

ConvTranspose2d

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

ConvTranspose3d

在一个input tensor上使用3D转置卷积。

LocalConv2d

对4D输入tensor进行通道分组并应用非统一kernel的空间卷积操作。

DeformableConv2d

可变形卷积

SlidingWindow

对输入tensor连续应用滑窗操作,每次拷贝滑窗内数据作为输出对应位置的数据。

SlidingWindowTranspose

滑窗操作的反向操作,对输入滑窗内数据求和作为输出。

池化层#

AvgPool2d

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

MaxPool2d

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

AdaptiveAvgPool2d

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

AdaptiveMaxPool2d

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

DeformablePSROIPooling

Padding层#

Pad

填充输入张量。

非线性激活层#

Sigmoid

对每个元素应用函数:

Softmax

应用softmax函数。

ReLU

对每个元素应用函数:

LeakyReLU

对每个元素应用函数:

PReLU

对每个元素应用函数:

SiLU

对每个元素应用函数:

GELU

对每个元素应用函数:

归一化层#

BatchNorm1d

Applies Batch Normalization over a 2D/3D tensor.

BatchNorm2d

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

SyncBatchNorm

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

GroupNorm

Simple implementation of GroupNorm.

InstanceNorm

Simple implementation of InstanceNorm.

LayerNorm

Simple implementation of LayerNorm.

LocalResponseNorm

对输入 Tensor 做 Local Response 正则化

Recurrent 层#

RNN

对输入序列应用一个多层 Elman RNN, 其中非线性单元采用的 \(\tanh\) 或者 \(\text{ReLU}\)

RNNCell

一个使用 tanh 或 ReLU 非线性的 Elman RNN 单元

LSTM

将多层 LSTM 应用于输入序列。

LSTMCell

长短期记忆 (LSTM) 单元。

线性变化层#

Identity

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

Linear

对输入应用线性变换。

Droupout 层#

Dropout

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

Sparse 层#

Embedding

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

Vision 层#

PixelShuffle

将形状为(, C x r^2, H, W)的向量中的元素重新排列为形状为(, C, H x r, W x r) 的向量,其中 r 是一个上采样因子,其中*是零或 batch 的维度。

融合算子#

ConvRelu2d

融合 Conv2drelu 运算的 Module.

ConvBn2d

融合 Conv2dBatchNorm2d 运算的 Module.

ConvBnRelu2d

融合 Conv2d, BatchNorm2drelu 运算的 Module.

BatchMatMulActivation

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

量化#

QuantStub

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

DequantStub

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

QAT Module#

容器#

QATModule

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

算子#

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.

ConvTranspose2d

A QATModule ConvTranspose2d with QAT support.

BatchMatMulActivation

支持 QAT 的 QATModule BatchMatMulActivation.

QuantStub

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

DequantStub

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

量化 Module#

QuantizedModule

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

算子#

Linear

量化(quantized)版本 Linear

Elemwise

量化(quantized)版本 Elemwise

Concat

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

Conv2d

量化(quantized)版本 Conv2d

ConvRelu2d

量化(quantized)版本 ConvRelu2d.

ConvBn2d

量化(quantized)版本 ConvBn2d

ConvBnRelu2d

量化(quantized)版本 ConvBnRelu2d.

ConvTranspose2d

ConvTranspose2d 的量化版本。

BatchMatMulActivation

BatchMatMulActivation 的量化版本。

QuantStub

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

DequantStub

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

External 层#

ExternOprSubgraph

导入一个序列化的 外部 Opr 子图。

TensorrtRuntimeSubgraph

导入一个序列化的 Tensorrt 运行时子图。

CambriconRuntimeSubgraph

导入一个序列化的 Cambricon 运行时子图。

AtlasRuntimeSubgraph

导入一个序列化的 Atlas 运行时子图。

MagicMindRuntimeSubgraph

导入一个序列化的 MagicMind 运行时子图。