main_quant.py

import os
import time
import numpy as np
import argparse
import functools
import shutil
import math
import multiprocessing

import paddle
import paddle.fluid as fluid
import reader
from mobilenet_ssd import mobile_net
from utility import add_arguments, print_arguments
from train import build_program
from train import train_parameters
from infer import draw_bounding_box_on_image

parser = argparse.ArgumentParser(description=__doc__)
add_arg = functools.partial(add_arguments, argparser=parser)
# yapf: disable
add_arg('learning_rate',    float, 0.0001,              "Learning rate.")
add_arg('batch_size',       int,   64,                  "Minibatch size.")
add_arg('epoc_num',         int,   20,                  "Epoch number.")
add_arg('use_gpu',          bool,  True,                "Whether use GPU.")
add_arg('parallel',         bool,  True,                "Whether train in parallel on multi-devices.")
add_arg('model_save_dir',   str,   'quant_model',       "The path to save model.")
add_arg('init_model',       str,   'ssd_mobilenet_v1_pascalvoc', "The init model path.")
add_arg('ap_version',       str,   '11point',           "mAP version can be integral or 11point.")
add_arg('image_shape',      str,   '3,300,300',         "Input image shape.")
add_arg('mean_BGR',         str,   '127.5,127.5,127.5', "Mean value for B,G,R channel which will be subtracted.")
add_arg('lr_epochs',        str,   '30,60',             "The learning decay steps.")
add_arg('lr_decay_rates',   str,   '1,0.1,0.01',        "The learning decay rates for each step.")
add_arg('data_dir',         str,   'data/pascalvoc',    "Data directory")
add_arg('act_quant_type',   str,   'abs_max',           "Quantize type of activation, whicn can be abs_max or range_abs_max")
add_arg('image_path',       str,   '',                  "The image used to inference and visualize.")
add_arg('confs_threshold',  float, 0.5,                 "Confidence threshold to draw bbox.")
add_arg('mode',             str,   'train',             "Job mode can be one of ['train', 'test', 'infer'].")
#yapf: enable

def test(exe, test_prog, map_eval, test_py_reader):
    _, accum_map = map_eval.get_map_var()
    map_eval.reset(exe)
    test_py_reader.start()
    try:
        batch = 0
        while True:
            test_map, = exe.run(test_prog, fetch_list=[accum_map])
            if batch % 10 == 0:
                print("Batch {0}, map {1}".format(batch, test_map))
            batch += 1
    except fluid.core.EOFException:
        test_py_reader.reset()
    finally:
        test_py_reader.reset()
    print("Test map {0}".format(test_map))
    return test_map


def save_model(exe, main_prog, model_save_dir, postfix):
    model_path = os.path.join(model_save_dir, postfix)
    if os.path.isdir(model_path):
        shutil.rmtree(model_path)
    fluid.io.save_persistables(exe, model_path, main_program=main_prog)


def train(args,
          data_args,
          train_params,
          train_file_list,
          val_file_list):

    model_save_dir = args.model_save_dir
    init_model = args.init_model
    epoc_num = args.epoc_num
    use_gpu = args.use_gpu
    parallel = args.parallel
    is_shuffle = True
    act_quant_type = args.act_quant_type

    if use_gpu:
        devices_num = fluid.core.get_cuda_device_count()
    else:
        devices_num = int(os.environ.get('CPU_NUM', multiprocessing.cpu_count()))

    batch_size = train_params['batch_size']
    batch_size_per_device = batch_size // devices_num
    num_workers = 4

    startup_prog = fluid.Program()
    train_prog = fluid.Program()
    test_prog = fluid.Program()

    train_py_reader, loss = build_program(
        main_prog=train_prog,
        startup_prog=startup_prog,
        train_params=train_params,
        is_train=True)
    test_py_reader, map_eval, _, _ = build_program(
        main_prog=test_prog,
        startup_prog=startup_prog,
        train_params=train_params,
        is_train=False)

    test_prog = test_prog.clone(for_test=True)

    transpiler = fluid.contrib.QuantizeTranspiler(weight_bits=8,
        activation_bits=8,
        activation_quantize_type=act_quant_type,
        weight_quantize_type='abs_max')

    transpiler.training_transpile(train_prog, startup_prog)
    transpiler.training_transpile(test_prog, startup_prog)

    place = fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace()
    exe = fluid.Executor(place)
    exe.run(startup_prog)

    if init_model:
        print('Load init model %s.' % init_model)
        def if_exist(var):
            return os.path.exists(os.path.join(init_model, var.name))
        fluid.io.load_vars(exe, init_model, main_program=train_prog,
                           predicate=if_exist)
    else:
        print('There is no init model.')

    if parallel:
        train_exe = fluid.ParallelExecutor(main_program=train_prog,
            use_cuda=True if use_gpu else False, loss_name=loss.name)

    train_reader = reader.train(data_args,
                                train_file_list,
                                batch_size_per_device,
                                shuffle=is_shuffle,
                                num_workers=num_workers)
    test_reader = reader.test(data_args, val_file_list, batch_size)
    train_py_reader.decorate_paddle_reader(train_reader)
    test_py_reader.decorate_paddle_reader(test_reader)

    train_py_reader.start()
    best_map = 0.
    for epoc in range(epoc_num):
        if epoc == 0:
            # test quantized model without quantization-aware training.
            test_map = test(exe, test_prog, map_eval, test_py_reader)
        batch = 0
        train_py_reader.start()
        while True:
            try:
                # train
                start_time = time.time()
                if parallel:
                    outs = train_exe.run(fetch_list=[loss.name])
                else:
                    outs = exe.run(train_prog, fetch_list=[loss])
                end_time = time.time()
                avg_loss = np.mean(np.array(outs[0]))
                if batch % 10 == 0:
                    print("Epoc {:d}, batch {:d}, loss {:.6f}, time {:.5f}".format(
                        epoc , batch, avg_loss, end_time - start_time))
            except (fluid.core.EOFException, StopIteration):
                train_reader().close()
                train_py_reader.reset()
                break
        test_map = test(exe, test_prog, map_eval, test_py_reader)
        save_model(exe, train_prog, model_save_dir, str(epoc))
        if test_map > best_map:
            best_map = test_map
            save_model(exe, train_prog, model_save_dir, 'best_map')
        print("Best test map {0}".format(best_map))


def eval(args, data_args, configs, val_file_list):
    init_model = args.init_model
    use_gpu = args.use_gpu
    act_quant_type = args.act_quant_type
    model_save_dir = args.model_save_dir

    batch_size = configs['batch_size']
    batch_size_per_device = batch_size

    startup_prog = fluid.Program()
    test_prog = fluid.Program()
    test_py_reader, map_eval, nmsed_out, image = build_program(
        main_prog=test_prog,
        startup_prog=startup_prog,
        train_params=configs,
        is_train=False)
    test_prog = test_prog.clone(for_test=True)

    transpiler = fluid.contrib.QuantizeTranspiler(weight_bits=8,
        activation_bits=8,
        activation_quantize_type=act_quant_type,
        weight_quantize_type='abs_max')
    transpiler.training_transpile(test_prog, startup_prog)

    place = fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace()
    exe = fluid.Executor(place)
    exe.run(startup_prog)

    def if_exist(var):
        return os.path.exists(os.path.join(init_model, var.name))
    fluid.io.load_vars(exe, init_model, main_program=test_prog,
                       predicate=if_exist)

    # freeze after load parameters
    transpiler.freeze_program(test_prog, place)

    test_reader = reader.test(data_args, val_file_list, batch_size)
    test_py_reader.decorate_paddle_reader(test_reader)

    test_map = test(exe, test_prog, map_eval, test_py_reader)
    print("Test model {0}, map {1}".format(init_model, test_map))
    # convert model to 8-bit before saving, but now Paddle can't load
    # the 8-bit model to do inference.
    # transpiler.convert_to_int8(test_prog, place)
    fluid.io.save_inference_model(model_save_dir, [image.name],
                                  [nmsed_out], exe, test_prog)


def infer(args, data_args):
    model_dir = args.init_model
    image_path = args.image_path
    confs_threshold = args.confs_threshold

    place = fluid.CUDAPlace(0) if args.use_gpu else fluid.CPUPlace()
    exe = fluid.Executor(place)
    [inference_program, feed , fetch] = fluid.io.load_inference_model(
        dirname=model_dir,
        executor=exe,
        model_filename='__model__')

    #print(np.array(fluid.global_scope().find_var('conv2d_20.w_0').get_tensor()))
    #print(np.max(np.array(fluid.global_scope().find_var('conv2d_20.w_0').get_tensor())))
    infer_reader = reader.infer(data_args, image_path)
    data = infer_reader()
    data = data.reshape((1,) + data.shape)
    outs = exe.run(inference_program,
                   feed={feed[0]: data},
                   fetch_list=fetch,
                   return_numpy=False)
    out = np.array(outs[0])
    draw_bounding_box_on_image(image_path, out, confs_threshold,
                               data_args.label_list)


if __name__ == '__main__':
    args = parser.parse_args()
    print_arguments(args)

    # for pascalvoc
    label_file = 'label_list'
    train_list = 'trainval.txt'
    val_list = 'test.txt'
    dataset = 'pascalvoc'

    mean_BGR = [float(m) for m in args.mean_BGR.split(",")]
    image_shape = [int(m) for m in args.image_shape.split(",")]
    lr_epochs = [int(m) for m in args.lr_epochs.split(",")]
    lr_rates = [float(m) for m in args.lr_decay_rates.split(",")]
    train_parameters[dataset]['image_shape'] = image_shape
    train_parameters[dataset]['batch_size'] = args.batch_size
    train_parameters[dataset]['lr'] = args.learning_rate
    train_parameters[dataset]['epoc_num'] = args.epoc_num
    train_parameters[dataset]['ap_version'] = args.ap_version
    train_parameters[dataset]['lr_epochs'] = lr_epochs
    train_parameters[dataset]['lr_decay'] = lr_rates

    data_args = reader.Settings(
        dataset=dataset,
        data_dir=args.data_dir,
        label_file=label_file,
        resize_h=image_shape[1],
        resize_w=image_shape[2],
        mean_value=mean_BGR,
        apply_distort=True,
        apply_expand=True,
        ap_version = args.ap_version)
    if args.mode == 'train':
        train(args, data_args, train_parameters[dataset], train_list, val_list)
    elif args.mode == 'test':
        eval(args, data_args, train_parameters[dataset], val_list)
    else:
        infer(args, data_args)