megengine.module.qat package

megengine.module.qat.batch_matmul_activation

class megengine.module.qat.batch_matmul_activation.BatchMatMulActivation(batch, in_features, out_features, bias=True, nonlinear_mode='IDENTITY', **kwargs)[source]

Bases: megengine.module.batch_matmul_activation.BatchMatMulActivation, megengine.module.qat.module.QATModule

forward(inp)[source]
classmethod from_float_module(float_module)[source]

Return a QATModule instance converted from a float Module instance.

megengine.module.qat.concat

class megengine.module.qat.concat.Concat[source]

Bases: megengine.module.concat.Concat, megengine.module.qat.module.QATModule

A QATModule to do functional concat with QAT support. Could be applied with Observer and FakeQuantize.

forward(inps, axis=0)[source]
classmethod from_float_module(float_module)[source]

Return a QATModule instance converted from a float Module instance.

megengine.module.qat.conv

class megengine.module.qat.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')[source]

Bases: megengine.module.conv.Conv2d, megengine.module.qat.module.QATModule

A QATModule Conv2d with QAT support. Could be applied with Observer and FakeQuantize.

calc_conv_qat(inp)[source]
forward(inp)[source]
classmethod from_float_module(float_module)[source]

Return a QATModule instance converted from a float Module instance.

class megengine.module.qat.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')[source]

Bases: megengine.module.qat.conv.Conv2d

A QATModule include Conv2d and Relu with QAT support. Could be applied with Observer and FakeQuantize.

forward(inp)[source]

megengine.module.qat.conv_bn

class megengine.module.qat.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)[source]

Bases: megengine.module.qat.conv_bn._ConvBnActivation2d

A fused QATModule including Conv2d, BatchNorm2d with QAT support. Could be applied with Observer and FakeQuantize.

forward(inp)[source]
class megengine.module.qat.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)[source]

Bases: megengine.module.qat.conv_bn._ConvBnActivation2d

A fused QATModule including Conv2d, BatchNorm2d and relu with QAT support. Could be applied with Observer and FakeQuantize.

forward(inp)[source]

megengine.module.qat.elemwise

class megengine.module.qat.elemwise.Elemwise(method)[source]

Bases: megengine.module.elemwise.Elemwise, megengine.module.qat.module.QATModule

A QATModule to do elemwise operator with QAT support. Could be applied with Observer and FakeQuantize.

Parameters

method – the elemwise method, see Elemwise for detail.

forward(*inps)[source]
classmethod from_float_module(float_module)[source]

Return a QATModule instance converted from a float Module instance.

with_weight = False

megengine.module.qat.linear

class megengine.module.qat.linear.Linear(in_features, out_features, bias=True, **kwargs)[source]

Bases: megengine.module.linear.Linear, megengine.module.qat.module.QATModule

A QATModule version of Linear. Could be applied with Observer and FakeQuantize.

Parameters
  • in_features (int) – size of each input sample.

  • out_features (int) – size of each output sample.

  • bias (bool) – If set to False, the layer will not learn an additive bias. Default: True

forward(x)[source]
classmethod from_float_module(float_module)[source]

Return a QATModule instance converted from a float Module instance.

megengine.module.qat.module

class megengine.module.qat.module.QATModule[source]

Bases: megengine.module.module.Module

Base class of quantized-float related Module, basically for QAT and Calibration.

Use from_float_module() to generate a instance from float Module. Or use quantize_qat() to do it recursively and automatically.

Can also be converted to QuantizedModule for deployment using quantize() further.

apply_quant_activation(target)[source]

Apply weight’s observer and fake_quant from qconfig on target.

apply_quant_weight(target)[source]

Apply weight’s observer and fake_quant from qconfig on target.

abstract classmethod from_float_module(float_module)[source]

Return a QATModule instance converted from a float Module instance.

get_activation_dtype()[source]

Get activation’s quantization dtype as the method from qconfig.

get_activation_qparams()[source]

Get activation’s quantization parameters.

get_weight_dtype()[source]

Get weight’s quantization dtype as the method from qconfig.

get_weight_qparams()[source]

Get weight’s quantization parameters.

set_fake_quant(enable)[source]
set_observer(enable)[source]
set_qconfig(qconfig)[source]

Set quantization related configs with qconfig, including observer and fake_quant for weight and activation.

with_act = True
with_weight = True

megengine.module.qat.quant_dequant

class megengine.module.qat.quant_dequant.DequantStub[source]

Bases: megengine.module.quant_dequant.DequantStub, megengine.module.qat.module.QATModule

A helper QATModule simply return input, but will de-quantize input after converted to QuantizedModule.

forward(inp)[source]
classmethod from_float_module(float_module)[source]

Return a QATModule instance converted from a float Module instance.

with_act = False
with_weight = False
class megengine.module.qat.quant_dequant.QuantStub[source]

Bases: megengine.module.quant_dequant.QuantStub, megengine.module.qat.module.QATModule

A helper QATModule simply return input, but will quantize input after converted to QuantizedModule.

forward(inp)[source]
classmethod from_float_module(float_module)[source]

Return a QATModule instance converted from a float Module instance.

with_weight = False