def from_float(cls, mod):
r"""Create a quantized sparse dynamic module from a float module.
We only care about the convert at this stage, no need for observers just yet.
"""
assert type(mod) == cls._FLOAT_MODULE, ' nnq.' + cls.__name__ + '.from_float only works for ' + \
cls._FLOAT_MODULE.__name__
# TODO: Need to add options to qconfig to avoid the calibration.
# TODO: Add calibration for the sparsity
assert hasattr(
mod, 'qconfig'), 'Input float module must have qconfig defined'
if type(mod) == nni.LinearReLU:
mod = mod[0]
if mod.qconfig is not None and mod.qconfig.weight is not None:
weight_observer = mod.qconfig.weight()
else:
# We have the circular import issues if we import the qconfig in the beginning of this file:
# https://github.com/pytorch/pytorch/pull/24231. The current workaround is to postpone the
# import until we need it.
from torch.ao.quantization.qconfig import default_dynamic_qconfig
weight_observer = default_dynamic_qconfig.weight()
# It is important to multiply by the mask BEFORE calling the `weight_observer`
# TODO (zaf): Mask might not be part of the qconfig (T83295194)
weight = mod.weight
if getattr(mod.qconfig, 'mask', False):
weight = mod.qconfig.mask * mod.weight
weight_observer(weight)
dtype = weight_observer.dtype
assert dtype == torch.qint8, 'Weight observer must have dtype torch.qint8'
w_sc, w_zp = weight_observer.calculate_qparams()
if isinstance(w_zp, torch.Tensor):
assert not torch.any(
w_zp.bool()), "All weight zero points must map to 0"
else:
assert w_zp == 0, 'Weight zero point must map to 0'
qweight = _quantize_weight(weight.float(), weight_observer)
row_block_size, col_block_size = LinearBlockSparsePattern.block_size()
qlinear = cls(mod.in_features,
mod.out_features,
row_block_size,
col_block_size,
dtype=dtype)
qlinear.set_weight_bias(qweight, mod.bias, row_block_size,
col_block_size)
return qlinear
def from_float(cls, mod):
r"""Create a quantized sparse module from a float module.
We only care about the convert at this stage, no need for observers just yet.
TODO: Need to figure out how to store the block shapes in the mod
"""
assert type(mod) == cls._FLOAT_MODULE, cls._get_name() + \
'.from_float only works for ' + cls._FLOAT_MODULE.__name__
# TODO: Need to add options to qconfig to avoid the calibration.
# TODO: Add calibration for the sparsity
assert hasattr(
mod, 'qconfig'), 'Input float module must have qconfig defined'
activation_post_process = mod.activation_post_process
weight_post_process = mod.qconfig.weight()
# Assumption is that the weight is already sparsified by the
# `sparsifier.convert`
weight = mod.weight
weight_post_process(weight)
dtype = weight_post_process.dtype
act_scale, act_zp = activation_post_process.calculate_qparams()
assert dtype == torch.qint8, 'Weight observer must have dtype torch.qint8'
w_sc, w_zp = weight_post_process.calculate_qparams()
if isinstance(w_zp, torch.Tensor):
assert not torch.any(
w_zp.bool()), "All weight zero points must map to 0"
else:
assert w_zp == 0, 'Weight zero point must map to 0'
qweight = _quantize_weight(weight.float(), weight_post_process)
row_block_size, col_block_size = LinearBlockSparsePattern.block_size()
qlinear = cls(mod.in_features,
mod.out_features,
row_block_size,
col_block_size,
dtype=dtype)
qlinear.set_weight_bias(qweight, mod.bias, row_block_size,
col_block_size)
qlinear.scale = float(act_scale)
qlinear.zero_point = int(act_zp)
return qlinear