vinx13/quantization_eval.py

## quantization_eval.py
import numpy as np
import logging
import argparse
import os
import mxnet as mx
from mxnet import gluon
from mxnet.gluon.model_zoo import vision
import tvm
import tvm.relay as relay
import tvm.relay.expr as _expr
import tvm.relay.transform as _transform
from tvm.contrib import graph_runtime
from scipy import stats
import pickle
import multiprocessing as mp


# Two functions for reading data from record file or raw images
def get_val_data(args,
                 rec_val,
                 batch_size,
                 num_workers=4,
                 shuffle=False):
    rec_val = os.path.expanduser(rec_val)
    mean_rgb = [123.68, 116.779, 103.939]
    std_rgb = [58.393, 57.12, 57.375]
    def batch_fn(batch, ctx):
        data = gluon.utils.split_and_load(batch.data[0], ctx_list=ctx, batch_axis=0)
        label = gluon.utils.split_and_load(batch.label[0], ctx_list=ctx, batch_axis=0)
        return data, label

    img_size = 299 if args.model == 'inceptionv3' else 224
    val_data = mx.io.ImageRecordIter(
        path_imgrec         = rec_val,
        preprocess_threads  = num_workers,
        shuffle             = shuffle,
        batch_size          = batch_size,
        resize              = 256,
        data_shape          = (3, img_size, img_size),
        mean_r              = mean_rgb[0],
        mean_g              = mean_rgb[1],
        mean_b              = mean_rgb[2],
        std_r               = std_rgb[0],
        std_g               = std_rgb[1],
        std_b               = std_rgb[2],
    )
    return val_data, batch_fn


def calibration_dataset():
    val_data, batch_fn = get_val_data(args, args.rec_val, args.batch_size, shuffle=True)
    val_data.reset()
    for i, batch in enumerate(val_data):
        if i*args.batch_size > args.calibration_samples:
            break
        data, label = batch_fn(batch, [mx.cpu(0)])
        yield {'data': data[0].asnumpy()}


def evaluate(args, graph, lib, params, ctx):
    """Evaluate on the validation set."""
    # setup dataset.
    batch_size = args.batch_size
    val_data, batch_fn = get_val_data(args, args.rec_val, batch_size)

    # create runtime module
    m = graph_runtime.create(graph, lib, ctx)
    m.set_input(**params)
    oshape = (batch_size, args.num_classes)
    out_arr = tvm.nd.empty(oshape, "float32")
    # setup evaluaiton metric
    acc_top1 = mx.metric.Accuracy()
    acc_top5 = mx.metric.TopKAccuracy(5)
    val_data.reset()
    acc_top1.reset()
    acc_top5.reset()

    # Execute
    for i, batch in enumerate(val_data):
        data, label = batch_fn(batch, [mx.cpu(0)])
        m.run(data=data[0].asnumpy())
        m.get_output(0, out_arr)
        acc_top1.update(label, [mx.nd.array(out_arr.asnumpy())])
        acc_top5.update(label, [mx.nd.array(out_arr.asnumpy())])

        if args.log_interval and not (i + 1) % args.log_interval:
            _, top1 = acc_top1.get()
            _, top5 = acc_top5.get()
            nsamples = (i + 1) * batch_size
            logging.info('[%d samples] validation: acc-top1=%f acc-top5=%f', nsamples, top1, top5)
    logging.info('[final] validation: acc-top1=%f acc-top5=%f', top1, top5)
    with open(args.record_file, "a") as f:
        f.write('{}, {} / {}\n'.format(
            args.model, top1, top5))


def calibrate_on_dataset(qgraph):
    profile_graph = relay.quantize.collect_stats(qgraph)
    with relay.build_config(opt_level=3):
        graph, lib, params = relay.build(relay.Module.from_expr(profile_graph), target=args.target)
    outputs = []
    m = graph_runtime.create(graph, lib, tvm.context(args.target, args.device_id))
    m.set_input(**params)

    num_outputs = m.get_num_outputs()
    outputs = [[] for i in range(num_outputs)]

    for batch_id, batch in enumerate(calibration_dataset()):
        print('batch {}..'.format(batch_id))
        m.set_input(**batch)
        m.run()
        for i in range(num_outputs):
            output = m.get_output(i).asnumpy()
            outputs[i].append(output)
    for i in range(num_outputs):
        outputs[i] = np.concatenate(outputs[i]).reshape(-1)
    with mp.Pool() as pool:
        scales = list(pool.map(relay.quantize.kl_divergence.kl_divergence_scale, outputs))
    return scales


def build_model(gluon_model, original):
    """Build with relay."""
    import tvm
    from tvm import relay
    from tvm.relay import quantize as qtz
    img_size = 299 if args.model == 'inceptionv3' else 224
    data_shape = (args.batch_size, 3, img_size, img_size)
    mod, params = relay.frontend.from_mxnet(gluon_model, {"data": data_shape})
    target = args.target
    ctx = tvm.context(target, args.device_id)

    if original:
        # run original model
        with relay.build_config(opt_level=3):
            graph, lib, params = relay.build(net, target, params=params)
        return graph, lib, params, ctx

    skip_conv_layers = [0]
    with relay.quantize.qconfig(store_lowbit_output=False, skip_conv_layers=skip_conv_layers):
        from tvm.relay.quantize.quantize import _bind_params
        graph = _bind_params(mod['main'], params)
        mod = relay.Module.from_expr(graph)
        optimize = _transform.Sequential([_transform.SimplifyInference(),
                                          _transform.FoldConstant(),
                                          _transform.FoldScaleAxis(),
                                          _transform.CanonicalizeOps(),
                                          _transform.FoldConstant()])
        with relay.build_config(opt_level=2):
            mod = optimize(mod)
            mod = relay.quantize.annotate()(mod)

            cache_file = '{}_scales.pkl'.format(args.model)
            if os.path.exists(cache_file):
                with open(cache_file, 'rb') as f:
                    scales = pickle.load(f)
            else:
                scales = calibrate_on_dataset(mod['main'])
                with open(cache_file, 'wb') as f:
                    pickle.dump(scales, f)
            if args.eval_power2:
                scales = list(map(lambda scale: 2**np.math.ceil(np.math.log(scale, 2)) if scale > 0 else 1.0, scales))
                weight_scales = 'power2'
            else:
                weight_scales = 'max'

            mod['main'] = relay.quantize.calibrate(mod['main'], weight_scales=weight_scales,
                                                   scales=scales)
            mod = relay.quantize.realize()(mod)


            mod = relay.transform.FoldConstant()(mod)
            graph, lib, params = relay.build(mod, target=args.target)

    return graph, lib, params, ctx


def save_model(name, graph, lib, params):
    with open(name + '.json', 'w') as f:
        f.write(graph)
    lib.export_library(name + '.so')
    with open(name + '.bin', 'wb') as f:
        f.write(relay.save_param_dict(params))


def main():
    gluon_model = vision.get_model(args.model, pretrained=True)
    graph, lib, params, ctx = build_model(gluon_model, args.original)
    logging.info("Finish building model %s...", args.model)
    evaluate(args, graph, lib, params, ctx)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Evaluate ImageNet validation accuracy")
    parser.add_argument("--rec-val", type=str, default="~/.mxnet/datasets/imagenet/rec/val.rec",
                        help="the validation data")
    parser.add_argument("--num-classes", type=int, default=1000,
                        help="batch size")
    parser.add_argument("--model", type=str, default="resnet50_v2",
                        help="Name of the model")
    parser.add_argument("--log-interval", type=int, default=100,
                        help="log interval")
    parser.add_argument("--batch-size", type=int, default=1,
                        help="batch size")
    parser.add_argument("--target", type=str, default="cuda",
                        help="target option")
    parser.add_argument("--original", action="store_true",
                        help='whether to use original graph')
    parser.add_argument('--save_model', type=str, default=None)
    parser.add_argument('--calibration_samples', type=int, default=100)
    parser.add_argument('--device-id', type=int, default=0)
    parser.add_argument('--eval-power2', action='store_true',
                        help='in this mode, scales are restricted to power-of-2 (weight: power2' \
                             'scale, activation: round kld to power2)')
    parser.add_argument('--record-file', type=str, default='record.csv',
                        help='file to save eval result')
    args = parser.parse_args()
    logging.basicConfig(level=logging.INFO)
    logging.info(args)
    main()
	import numpy as np
	import logging
	import argparse
	import os
	import mxnet as mx
	from mxnet import gluon
	from mxnet.gluon.model_zoo import vision
	import tvm
	import tvm.relay as relay
	import tvm.relay.expr as _expr
	import tvm.relay.transform as _transform
	from tvm.contrib import graph_runtime
	from scipy import stats
	import pickle
	import multiprocessing as mp


	# Two functions for reading data from record file or raw images
	def get_val_data(args,
	rec_val,
	batch_size,
	num_workers=4,
	shuffle=False):
	rec_val = os.path.expanduser(rec_val)
	mean_rgb = [123.68, 116.779, 103.939]
	std_rgb = [58.393, 57.12, 57.375]
	def batch_fn(batch, ctx):
	data = gluon.utils.split_and_load(batch.data[0], ctx_list=ctx, batch_axis=0)
	label = gluon.utils.split_and_load(batch.label[0], ctx_list=ctx, batch_axis=0)
	return data, label

	img_size = 299 if args.model == 'inceptionv3' else 224
	val_data = mx.io.ImageRecordIter(
	path_imgrec = rec_val,
	preprocess_threads = num_workers,
	shuffle = shuffle,
	batch_size = batch_size,
	resize = 256,
	data_shape = (3, img_size, img_size),
	mean_r = mean_rgb[0],
	mean_g = mean_rgb[1],
	mean_b = mean_rgb[2],
	std_r = std_rgb[0],
	std_g = std_rgb[1],
	std_b = std_rgb[2],
	)
	return val_data, batch_fn


	def calibration_dataset():
	val_data, batch_fn = get_val_data(args, args.rec_val, args.batch_size, shuffle=True)
	val_data.reset()
	for i, batch in enumerate(val_data):
	if i*args.batch_size > args.calibration_samples:
	break
	data, label = batch_fn(batch, [mx.cpu(0)])
	yield {'data': data[0].asnumpy()}


	def evaluate(args, graph, lib, params, ctx):
	"""Evaluate on the validation set."""
	# setup dataset.
	batch_size = args.batch_size
	val_data, batch_fn = get_val_data(args, args.rec_val, batch_size)

	# create runtime module
	m = graph_runtime.create(graph, lib, ctx)
	m.set_input(**params)
	oshape = (batch_size, args.num_classes)
	out_arr = tvm.nd.empty(oshape, "float32")
	# setup evaluaiton metric
	acc_top1 = mx.metric.Accuracy()
	acc_top5 = mx.metric.TopKAccuracy(5)
	val_data.reset()
	acc_top1.reset()
	acc_top5.reset()

	# Execute
	for i, batch in enumerate(val_data):
	data, label = batch_fn(batch, [mx.cpu(0)])
	m.run(data=data[0].asnumpy())
	m.get_output(0, out_arr)
	acc_top1.update(label, [mx.nd.array(out_arr.asnumpy())])
	acc_top5.update(label, [mx.nd.array(out_arr.asnumpy())])

	if args.log_interval and not (i + 1) % args.log_interval:
	_, top1 = acc_top1.get()
	_, top5 = acc_top5.get()
	nsamples = (i + 1) * batch_size
	logging.info('[%d samples] validation: acc-top1=%f acc-top5=%f', nsamples, top1, top5)
	logging.info('[final] validation: acc-top1=%f acc-top5=%f', top1, top5)
	with open(args.record_file, "a") as f:
	f.write('{}, {} / {}\n'.format(
	args.model, top1, top5))


	def calibrate_on_dataset(qgraph):
	profile_graph = relay.quantize.collect_stats(qgraph)
	with relay.build_config(opt_level=3):
	graph, lib, params = relay.build(relay.Module.from_expr(profile_graph), target=args.target)
	outputs = []
	m = graph_runtime.create(graph, lib, tvm.context(args.target, args.device_id))
	m.set_input(**params)

	num_outputs = m.get_num_outputs()
	outputs = [[] for i in range(num_outputs)]

	for batch_id, batch in enumerate(calibration_dataset()):
	print('batch {}..'.format(batch_id))
	m.set_input(**batch)
	m.run()
	for i in range(num_outputs):
	output = m.get_output(i).asnumpy()
	outputs[i].append(output)
	for i in range(num_outputs):
	outputs[i] = np.concatenate(outputs[i]).reshape(-1)
	with mp.Pool() as pool:
	scales = list(pool.map(relay.quantize.kl_divergence.kl_divergence_scale, outputs))
	return scales


	def build_model(gluon_model, original):
	"""Build with relay."""
	import tvm
	from tvm import relay
	from tvm.relay import quantize as qtz
	img_size = 299 if args.model == 'inceptionv3' else 224
	data_shape = (args.batch_size, 3, img_size, img_size)
	mod, params = relay.frontend.from_mxnet(gluon_model, {"data": data_shape})
	target = args.target
	ctx = tvm.context(target, args.device_id)

	if original:
	# run original model
	with relay.build_config(opt_level=3):
	graph, lib, params = relay.build(net, target, params=params)
	return graph, lib, params, ctx

	skip_conv_layers = [0]
	with relay.quantize.qconfig(store_lowbit_output=False, skip_conv_layers=skip_conv_layers):
	from tvm.relay.quantize.quantize import _bind_params
	graph = _bind_params(mod['main'], params)
	mod = relay.Module.from_expr(graph)
	optimize = _transform.Sequential([_transform.SimplifyInference(),
	_transform.FoldConstant(),
	_transform.FoldScaleAxis(),
	_transform.CanonicalizeOps(),
	_transform.FoldConstant()])
	with relay.build_config(opt_level=2):
	mod = optimize(mod)
	mod = relay.quantize.annotate()(mod)

	cache_file = '{}_scales.pkl'.format(args.model)
	if os.path.exists(cache_file):
	with open(cache_file, 'rb') as f:
	scales = pickle.load(f)
	else:
	scales = calibrate_on_dataset(mod['main'])
	with open(cache_file, 'wb') as f:
	pickle.dump(scales, f)
	if args.eval_power2:
	scales = list(map(lambda scale: 2**np.math.ceil(np.math.log(scale, 2)) if scale > 0 else 1.0, scales))
	weight_scales = 'power2'
	else:
	weight_scales = 'max'

	mod['main'] = relay.quantize.calibrate(mod['main'], weight_scales=weight_scales,
	scales=scales)
	mod = relay.quantize.realize()(mod)


	mod = relay.transform.FoldConstant()(mod)
	graph, lib, params = relay.build(mod, target=args.target)

	return graph, lib, params, ctx


	def save_model(name, graph, lib, params):
	with open(name + '.json', 'w') as f:
	f.write(graph)
	lib.export_library(name + '.so')
	with open(name + '.bin', 'wb') as f:
	f.write(relay.save_param_dict(params))


	def main():
	gluon_model = vision.get_model(args.model, pretrained=True)
	graph, lib, params, ctx = build_model(gluon_model, args.original)
	logging.info("Finish building model %s...", args.model)
	evaluate(args, graph, lib, params, ctx)


	if __name__ == "__main__":
	parser = argparse.ArgumentParser(description="Evaluate ImageNet validation accuracy")
	parser.add_argument("--rec-val", type=str, default="~/.mxnet/datasets/imagenet/rec/val.rec",
	help="the validation data")
	parser.add_argument("--num-classes", type=int, default=1000,
	help="batch size")
	parser.add_argument("--model", type=str, default="resnet50_v2",
	help="Name of the model")
	parser.add_argument("--log-interval", type=int, default=100,
	help="log interval")
	parser.add_argument("--batch-size", type=int, default=1,
	help="batch size")
	parser.add_argument("--target", type=str, default="cuda",
	help="target option")
	parser.add_argument("--original", action="store_true",
	help='whether to use original graph')
	parser.add_argument('--save_model', type=str, default=None)
	parser.add_argument('--calibration_samples', type=int, default=100)
	parser.add_argument('--device-id', type=int, default=0)
	parser.add_argument('--eval-power2', action='store_true',
	help='in this mode, scales are restricted to power-of-2 (weight: power2' \
	'scale, activation: round kld to power2)')
	parser.add_argument('--record-file', type=str, default='record.csv',
	help='file to save eval result')
	args = parser.parse_args()
	logging.basicConfig(level=logging.INFO)
	logging.info(args)
	main()