import mxnet as mx
from mxnet.gluon.model_zoo import vision
import psutil
import os

def get_process_memory():
    process = psutil.Process(os.getpid())
    mem_info = process.memory_info()
    return mem_info.rss * 1e-6


batch_shape = (1, 3, 224, 224)
data = mx.np.random.normal(size=batch_shape)

print("memory before loading model: ", get_process_memory())
net = vision.resnet50_v1(pretrained=True)
print("memory after loading model: ", get_process_memory())
out = net(data)
out.wait_to_read()
print("memory after fp32 forward pass", get_process_memory())

dataset = mx.gluon.data.ArrayDataset(data)
data_loader = mx.gluon.data.DataLoader(dataset, batch_size=1)
net_quantized = mx.contrib.quant.quantize_net(net, quantized_dtype='int8',
                                                quantize_mode="smart",
                                                calib_mode='naive',
                                                calib_data=data_loader,
                                                num_calib_batches=1,
                                                ctx=mx.current_context())

print("memory after quantization: ", get_process_memory())

outputs = net_quantized(data)
outputs.wait_to_read()
print("memory after int8 forward pass: ", get_process_memory())

memory before loading model:  213.430272
[15:14:11] ../src/storage/storage.cc:202: Using Pooled (Naive) StorageManager for CPU
memory after loading model:  530.702336
memory after fp32 forward pass 611.241984
/home/bg/work/MXNet/python/mxnet/gluon/block.py:1918: UserWarning: Cannot decide type for the following arguments. Consider providing them as input:
        data: None
  input_sym_arg_type = in_param.infer_type()[0]
/home/bg/work/MXNet/python/mxnet/gluon/block.py:1251: UserWarning: register_op_hook is experimental when static_alloc=True / static_shape=True  and may not work correctly
  warnings.warn("register_op_hook is experimental when static_alloc=True / static_shape=True "
memory after quantization:  1064.57088
memory after int8 forward pass:  1071.005696

memory before loading model:  214.28633599999998
[15:13:17] ../src/storage/storage.cc:202: Using Pooled (Naive) StorageManager for CPU
memory after loading model:  531.2593919999999
memory after fp32 forward pass 609.513472
/home/bg/work/MXNet/python/mxnet/gluon/block.py:1918: UserWarning: Cannot decide type for the following arguments. Consider providing them as input:
        data: None
  input_sym_arg_type = in_param.infer_type()[0]
/home/bg/work/MXNet/python/mxnet/gluon/block.py:1251: UserWarning: register_op_hook is experimental when static_alloc=True / static_shape=True  and may not work correctly
  warnings.warn("register_op_hook is experimental when static_alloc=True / static_shape=True "
memory after quantization:  890.273792
memory after int8 forward pass:  895.2258559999999