def quantize(self, tune_cfg, model, dataloader, q_func=None): """The function is used to do MXNet calibration and quanitization in post-training quantization. Args: tune_cfg (dict): quantization config. model (object): model need to do quantization. dataloader (object): calibration dataset. q_func (optional): training function for quantization aware training mode, unimplement yet for MXNet. Returns: (dict): quantized model """ assert q_func is None, "quantization aware training mode is not support on mxnet" # get symbol from FP32 model if isinstance(model.model, mx.gluon.HybridBlock): # transfer hybridblock to symbol sym, arg_params, aux_params, calib_data = \ self._get_gluon_symbol(model.model, dataloader=dataloader) data_names = [pair[0] for pair in calib_data.provide_data] self.__config_dict['calib_data'] = calib_data elif isinstance(model.model[0], mx.symbol.Symbol): sym, arg_params, aux_params = model.model self.__config_dict['calib_data'] = dataloader else: raise ValueError( 'Need a symbol model or HybridBlock model, while received %s' % str( type(model.model))) self._cfg_to_qconfig(tune_cfg) self.th_dict = None qconfig = self.__config_dict sym = self._get_backedn_graph(sym, qconfig['ctx']) # 1. quantize_symbol if _check_version(mx.__version__, '1.7.0'): qsym, calib_layer = mx.contrib.quantization._quantize_symbol( sym, qconfig['ctx'], excluded_symbols=qconfig['excluded_sym_names'], excluded_operators=qconfig['excluded_op_names'], offline_params=list(arg_params.keys()), quantized_dtype=qconfig['quantized_dtype'], quantize_mode=qconfig['quantize_mode'], quantize_granularity=qconfig['quantize_granularity']) else: qsym, calib_layer = mx.contrib.quantization._quantize_symbol( sym, qconfig['ctx'], excluded_symbols=qconfig['excluded_sym_names'], excluded_operators=qconfig['excluded_op_names'], offline_params=list(arg_params.keys()), quantized_dtype=qconfig['quantized_dtype'], quantize_mode=qconfig['quantize_mode'],) # 2. Do calibration to get th_dict th_dict = self._get_calibration_th( sym, arg_params, aux_params, calib_layer, qconfig) self.th_dict = th_dict # 3. set th_dict to quantized symbol qsym = mx.contrib.quantization._calibrate_quantized_sym(qsym, th_dict) # 4. quantize params qarg_params = mx.contrib.quantization._quantize_params( qsym, arg_params, th_dict) qsym = self._get_backedn_graph(qsym, qconfig['ctx']) if isinstance(model.model, mx.gluon.HybridBlock): from mxnet.gluon import SymbolBlock data_sym = [] for name in data_names: data_sym.append(mx.sym.var(name)) net = SymbolBlock(qsym, data_sym) with TemporaryDirectory() as tmpdirname: prefix = os.path.join(tmpdirname, 'tmp') param_name = '%s-%04d.params' % (prefix + 'net-quantized', 0) save_dict = {('arg:%s' % k): v.as_in_context(mx.cpu()) for k, v in qarg_params.items()} save_dict.update({('aux:%s' % k): v.as_in_context(mx.cpu()) for k, v in aux_params.items()}) mx.ndarray.save(param_name, save_dict) # pylint: disable=no-member net.collect_params().load(param_name, cast_dtype=True, dtype_source='saved') net.collect_params().reset_ctx(self.__config_dict['ctx']) return Model(net) return Model((qsym, qarg_params, aux_params))
class InceptionBNStages(HybridBlock): def __init__(self, **kwargs): super(InceptionBNStages, self).__init__(**kwargs) # data data = mx.symbol.Variable(name="data") # stage 1 conv1 = get_conv(data=data, num_filter=64, kernel=(7, 7), stride=(2, 2), pad=(3, 3), name='1') pool1 = mx.symbol.Pooling(data=conv1, kernel=(3, 3), stride=(2, 2), name='pool_1', pool_type='max') conv2red = get_conv(data=pool1, num_filter=64, kernel=(1, 1), stride=(1, 1), name='2_red') conv2 = get_conv(data=conv2red, num_filter=192, kernel=(3, 3), stride=(1, 1), pad=(1, 1), name='2') # stage 2 pool2 = mx.symbol.Pooling(data=data, kernel=(3, 3), stride=(2, 2), name='pool_2', pool_type='max') in3a = get_inception_a(pool2, 64, 64, 64, 64, 96, "avg", 32, '3a') in3b = get_inception_a(in3a, 64, 64, 96, 64, 96, "avg", 64, '3b') in3c = get_inception_b(in3b, 128, 160, 64, 96, '3c') # stage 3 in4a = get_inception_a(data, 224, 64, 96, 96, 128, "avg", 128, '4a') in4b = get_inception_a(in4a, 192, 96, 128, 96, 128, "avg", 128, '4b') in4c = get_inception_a(in4b, 160, 128, 160, 128, 160, "avg", 128, '4c') in4d = get_inception_a(in4c, 96, 128, 192, 160, 192, "avg", 128, '4d') in4e = get_inception_b(in4d, 128, 192, 192, 256, '4e') # stage 4 in5a = get_inception_a(data, 352, 192, 320, 160, 224, "avg", 128, '5a') in5b = get_inception_a(in5a, 352, 192, 320, 192, 224, "max", 128, '5b') # stage 5 pool = mx.symbol.Pooling(data=data, kernel=(7, 7), stride=(1, 1), name="global_pool", pool_type='avg') flatten = mx.symbol.Flatten(data=pool) self._stage_1 = SymbolBlock(outputs=conv2, inputs=mx.sym.var('data')) self._stage_2 = SymbolBlock(outputs=in3c, inputs=mx.sym.var('data')) self._stage_3 = SymbolBlock(outputs=in4e, inputs=mx.sym.var('data')) self._stage_4 = SymbolBlock(outputs=in5b, inputs=mx.sym.var('data')) self._stage_5 = SymbolBlock(outputs=flatten, inputs=mx.sym.var('data')) def hybrid_forward(self, F, x): stage1 = self._stage_1(x) stage2 = self._stage_2(stage1) stage3 = self._stage_3(stage2) stage4 = self._stage_4(stage3) stage5 = self._stage_5(stage4) return stage1, stage2, stage3, stage4, stage5 def load(self, param_file, ctx=mx.cpu()): self._stage_1.collect_params().load(param_file, ctx=ctx, ignore_extra=True) self._stage_2.collect_params().load(param_file, ctx=ctx, ignore_extra=True) self._stage_3.collect_params().load(param_file, ctx=ctx, ignore_extra=True) self._stage_4.collect_params().load(param_file, ctx=ctx, ignore_extra=True) self._stage_5.collect_params().load(param_file, ctx=ctx, ignore_extra=True)