def _quantize(cls,
sym: Symbol,
sym_ctrl: SymbolStats,
qrec: FloatQRec = None,
**kwargs) -> Tuple[Symbol, FloatQRec]:
prequantized_variables = kwargs.get('prequantized_variables', {})
qtypes = kwargs.get('qtypes', {})
# if the symbol is in prequantized_variables it is probably
# an output from a previous function so already has a quant record
if sym.name in prequantized_variables:
sym.qrec = prequantized_variables[sym.name]
return (sym, prequantized_variables[sym.name])
max_val = sym_ctrl.get_max(sym)
# see if an nntool quantizer qtype is available
if not qrec and qtypes and sym.name in qtypes:
qrec = cls.qrec_from_qtype(qtypes[sym.name], max_val)
if qrec:
sym.qrec = qrec
return (sym, qrec)
out_dtype = kwargs.get('out_dtype', np.float32)
return FloatQRec(dtype=out_dtype, max_val=max_val, min_val=-max_val)
def _quantize(cls,
sym: Symbol,
sym_ctrl: SymbolStats,
qrec: FloatQRec = None,
**kwargs) -> Tuple[Symbol, FloatQRec]:
in_syms, in_qrecs = zip(*[
cls.quantize(inner_sym, sym_ctrl, **kwargs)
for inner_sym in sym.contents
])
in_syms, in_qrecs = cls.cast_symbols(in_syms,
in_qrecs,
dtype=np.float32)
max_val = sym_ctrl.get_max(sym)
min_val = None if max_val is None else -max_val
out_qrec = FloatQRec(dtype=np.float32,
min_val=min_val,
max_val=max_val)
if isinstance(sym, Cos):
qsym = FastFloatCos
else:
qsym = FastFloatSin
return (qsym(*in_syms), out_qrec)
def _quantize(cls,
sym: Symbol,
sym_ctrl: SymbolStats,
qrec: FloatQRec = None,
**kwargs) -> Tuple[Symbol, FloatQRec]:
return sym, FloatQRec(dtype=np.float32,
min_val=sym.value[0],
max_val=sym.value[0])
def _quantize(cls,
sym: Symbol,
sym_ctrl: SymbolStats,
qrec: FloatQRec = None,
**kwargs) -> Tuple[Symbol, FloatQRec]:
in_syms, in_qrecs = zip(*[
cls.quantize(inner_sym, sym_ctrl, **kwargs)
for inner_sym in sym.contents
])
in_syms, in_qrecs = cls.cast_symbols(in_syms,
in_qrecs,
dtype=np.float32)
max_val = sym_ctrl.get_max(sym)
min_val = None if max_val is None else -max_val
out_qrec = FloatQRec(dtype=np.float32,
min_val=min_val,
max_val=max_val)
if isinstance(sym, Cos):
qsym = FastFloatCos
elif isinstance(sym, RSqrt):
qsym = FastFloatRSqrt
elif isinstance(sym, Sin):
qsym = FastFloatSin
elif isinstance(sym, Sigmoid):
qsym = FastFloatSigmoid
elif isinstance(sym, TanH):
qsym = FastFloatTanH
elif isinstance(sym, Pow):
base_val = in_syms[0].resolve()
if isinstance(base_val, Constant) and base_val.value == 2:
in_syms = [in_syms[1]]
qsym = FastFloatPow2
else:
qsym = FastFloatPow
elif isinstance(sym, Log):
qsym = FastFloatLog
elif isinstance(sym, Exp):
qsym = FastFloatExp
elif isinstance(sym, Abs):
if sym.dtype == np.float32:
qsym = Abs
else:
qsym = AbsF
return (qsym(*in_syms), out_qrec)
def _quantize(cls,
sym: Symbol,
sym_ctrl: SymbolStats,
qrec: FloatQRec = None,
**kwargs) -> Tuple[Symbol, FloatQRec]:
in_syms, in_qrecs = zip(*[
cls.quantize(inner_sym, sym_ctrl, **kwargs)
for inner_sym in sym.contents
])
sym_cls = sym.__class__
in_syms, in_qrecs = cls.cast_symbols(in_syms, in_qrecs)
max_val = sym_ctrl.get_max(sym)
min_val = None if max_val is None else -max_val
out_qrec = FloatQRec(dtype=np.float32,
min_val=min_val,
max_val=max_val)
return (Cast(sym_cls(*in_syms), dtype=np.float32), out_qrec)
def qrec_from_qtype(cls, qtype, max_val):
if qtype.dtype == np.int8 or qtype.dtype == np.int16:
if qtype.dtype == np.int8:
if len(qtype.scale) > 1:
qtype.scale = np.max(qtype.scale)
q = 7
dtype = np.int8
elif qtype.dtype == np.int16:
if len(qtype.scale) > 1:
qtype.scale = np.max(qtype.scale)
q = 15
dtype = np.int16
return Q15ScaleQRec(dtype,
max_val,
q,
max_val=max_val,
min_val=-max_val)
elif qtype.dtype == np.float32 or qtype.dtype == bfloat16:
return FloatQRec(dtype=qtype.dtype,
max_val=max_val,
min_val=-max_val)
else:
return None
def cast_symbols(in_syms, qrecs, dtype=np.float32):
return zip(*[(Cast(sym, dtype=dtype),
FloatQRec.override(qrec, dtype=dtype)
) if qrec.dtype != dtype else (sym, qrec)
for sym, qrec in zip(in_syms, qrecs)])