train_imagenet.py

import os
import argparse
import logging
logging.basicConfig(level=logging.DEBUG)
from common import find_mxnet, data, fit
from common.util import download_file
import mxnet as mx


def add_binary_args(parser):
    parser.add_argument('--bits-w', type=int, default=1,
                       help='number of bits for weights')
    parser.add_argument('--bits-a', type=int, default=1,
                       help='number of bits for activations')
    parser.add_argument('--reduction', type=float, default=1,
                       help='reduction in transition blocks (densenet only)')
    parser.add_argument('--growth-rate', type=int, default=32,
                       help='number of initial filters (densenet only)')

if __name__ == '__main__':

    # parse args
    parser = argparse.ArgumentParser(description="train imagenet1K",
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    fit.add_fit_args(parser)
    data.add_data_args(parser)
    data.add_data_aug_args(parser)
    # use a large aug level
    data.set_data_aug_level(parser, 3)

    parser.add_argument('--pretrained', type=str,
                    help='the pre-trained model')

    parser.add_argument('--log', dest='log_file', type=str, default="train.log",
                    help='save training log to file')

    parser.add_argument('--visualize', action='store_true',
                    help='visualize network instead of train it')

    add_binary_args(parser)

    parser.set_defaults(
        # network
        network        = 'resnet',
        num_layers     = 18,

        # data
        num_classes      = 1000,
        num_examples     = 1281167,
        image_shape      = '3,224,224',
        min_random_scale = 1, # if input image has min size k, suggest to use
                              # 256.0/x, e.g. 0.533 for 480
        # train
        num_epochs       = 40,
        lr_step_epochs   = '10,20,30',
        lr               = 0.01,
        lr_factor        = 0.1,
        batch_size       = 32,
        optimizer        = 'Nadam',
        disp_batches     = 10,
        top_k            = 5,
        data_train       = '/data/imagenet1k/imagenet1k-train.rec',
        data_val         = '/data/imagenet1k/imagenet1k-val.rec'
    )
    args = parser.parse_args()

    # set up logger
    log_file = args.log_file
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    if log_file:
        fh = logging.FileHandler(log_file)
        logger.addHandler(fh)

    devs = mx.cpu() if args.gpus is None or args.gpus is '' else [
    mx.gpu(int(i)) for i in args.gpus.split(',')]

    # load network
    from importlib import import_module
    net = import_module('symbols.'+args.network)
    sym = net.get_symbol(**vars(args))

    if args.visualize:
        mx.visualization.print_summary(sym, shape={'data': (1, 3, 224, 224)})
        # digraph = mx.visualization.plot_network(sym, save_format='jpg')
        # digraph.render()
    else:
        #load pretrained
        args_params=None
        auxs_params=None

        # train
        if args_params and auxs_params:
            fit.fit(
                args,
                sym,
                data.get_rec_iter,
                arg_params=args_params,
                aux_params=auxs_params)
        else:
            fit.fit(
                args,
                sym,
                data.get_rec_iter)