parser = argparse.ArgumentParser(description='Train R-FCN network')

# general

parser.add_argument('--cfg', help='experiment configure file name', required=True, type=str)

args, rest = parser.parse_known_args()

# update config

update_config(args.cfg)

# training

parser.add_argument('--frequent', help='frequency of logging', default=config.default.frequent, type=int)

args = parser.parse_args()

logger, final_output_path = create_logger(config.output_path, args.cfg, config.dataset.image_set)

prefix = os.path.join(final_output_path, prefix)

# load symbol

shutil.copy2(os.path.join(curr_path, 'symbols', config.symbol + '.py'), final_output_path)

sym_instance = eval(config.symbol + '.' + config.symbol)()

sym = sym_instance.get_symbol(config, is_train=True)

feat_sym = sym.get_internals()['rpn_cls_score_output']

# setup multi-gpu

batch_size = len(ctx)

input_batch_size = config.TRAIN.BATCH_IMAGES * batch_size

# print config

pprint.pprint(config)

logger.info('training config:{}\n'.format(pprint.pformat(config)))

# load dataset and prepare imdb for training

image_sets = [iset for iset in config.dataset.image_set.split('+')]

roidbs = [load_gt_roidb(config.dataset.dataset, image_set, config.dataset.root_path, config.dataset.dataset_path,

flip=config.TRAIN.FLIP)

for image_set in image_sets]

roidb = merge_roidb(roidbs)

roidb = filter_roidb(roidb, config)

# load training data

train_data = AnchorLoader(feat_sym, roidb, config, batch_size=input_batch_size, shuffle=config.TRAIN.SHUFFLE, ctx=ctx,

feat_stride=config.network.RPN_FEAT_STRIDE, anchor_scales=config.network.ANCHOR_SCALES,

anchor_ratios=config.network.ANCHOR_RATIOS, aspect_grouping=config.TRAIN.ASPECT_GROUPING)

# infer max shape

max_data_shape = [('data', (config.TRAIN.BATCH_IMAGES, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES])))]

max_data_shape, max_label_shape = train_data.infer_shape(max_data_shape)

max_data_shape.append(('gt_boxes', (config.TRAIN.BATCH_IMAGES, 100, 5)))

print('providing maximum shape', max_data_shape, max_label_shape)

data_shape_dict = dict(train_data.provide_data_single + train_data.provide_label_single)

pprint.pprint(data_shape_dict)

sym_instance.infer_shape(data_shape_dict)

# load and initialize params

#if config.TRAIN.RESUME:

# print('continue training from ', begin_epoch)

# arg_params, aux_params = load_param(prefix, begin_epoch, convert=True)

#else:

# arg_params, aux_params = load_param(pretrained, epoch, convert=True)

# sym_instance.init_weight(config, arg_params, aux_params)

print('transfer learning...')

# Choose the initialization weights (COCO or UADETRAC or pretrained)

#arg_params, aux_params = load_param('/raid10/home_ext/Deformable-ConvNets/output/rfcn_dcn_Shuo_UADTRAC/resnet_v1_101_voc0712_rfcn_dcn_Shuo_UADETRAC/trainlist_full/rfcn_UADTRAC', 5, convert=True)

#arg_params, aux_params = load_param('/raid10/home_ext/Deformable-ConvNets/model/rfcn_dcn_coco', 0, convert=True)

arg_params, aux_params = load_param('/raid10/home_ext/Deformable-ConvNets/output/rfcn_dcn_Shuo_AICity/resnet_v1_101_voc0712_rfcn_dcn_Shuo_AICityVOC1080_FreezeCOCO_rpnOnly_all/1080_all/rfcn_AICityVOC1080_FreezeCOCO_rpnOnly_all', 4, convert=True)

sym_instance.init_weight_Shuo(config, arg_params, aux_params)

# check parameter shapes

sym_instance.check_parameter_shapes(arg_params, aux_params, data_shape_dict)

# create solver

fixed_param_prefix = config.network.FIXED_PARAMS

data_names = [k[0] for k in train_data.provide_data_single]

label_names = [k[0] for k in train_data.provide_label_single]

mod = MutableModule(sym, data_names=data_names, label_names=label_names,

logger=logger, context=ctx, max_data_shapes=[max_data_shape for _ in range(batch_size)],

max_label_shapes=[max_label_shape for _ in range(batch_size)], fixed_param_prefix=fixed_param_prefix)

#freeze parameters using fixed_param_names:list of str

para_file = open('/raid10/home_ext/Deformable-ConvNets/rfcn/symbols/arg_params.txt')

para_list = [line.split('<')[0] for line in para_file.readlines()]

# para_list.remove('rfcn_cls_weight')

# para_list.remove('rfcn_cls_bias')

# para_list.remove('rfcn_cls_offset_t_weight')

# para_list.remove('rfcn_cls_offset_t_bias')

#

para_list.remove('res5a_branch2b_offset_weight')

para_list.remove('res5a_branch2b_offset_bias')

para_list.remove('res5b_branch2b_offset_weight')

para_list.remove('res5b_branch2b_offset_bias')

para_list.remove('res5c_branch2b_offset_weight')

para_list.remove('res5c_branch2b_offset_bias')

para_list.remove('conv_new_1_weight')

para_list.remove('conv_new_1_bias')

para_list.remove('rfcn_bbox_weight')

para_list.remove('rfcn_bbox_bias')

para_list.remove('rfcn_bbox_offset_t_weight')

para_list.remove('rfcn_bbox_offset_t_bias')

mod = MutableModule_Shuo(sym, data_names=data_names, label_names=label_names,

logger=logger, context=ctx, max_data_shapes=[max_data_shape for _ in range(batch_size)],

max_label_shapes=[max_label_shape for _ in range(batch_size)], fixed_param_prefix=fixed_param_prefix,

fixed_param_names=para_list)

if config.TRAIN.RESUME:

mod._preload_opt_states = '%s-%04d.states'%(prefix, begin_epoch)

# decide training params

# metric

rpn_eval_metric = metric.RPNAccMetric()

rpn_cls_metric = metric.RPNLogLossMetric()

rpn_bbox_metric = metric.RPNL1LossMetric()

eval_metric = metric.RCNNAccMetric(config)

cls_metric = metric.RCNNLogLossMetric(config)

bbox_metric = metric.RCNNL1LossMetric(config)

eval_metrics = mx.metric.CompositeEvalMetric()

# rpn_eval_metric, rpn_cls_metric, rpn_bbox_metric, eval_metric, cls_metric, bbox_metric

for child_metric in [rpn_eval_metric, rpn_cls_metric, rpn_bbox_metric, eval_metric, cls_metric, bbox_metric]:

eval_metrics.add(child_metric)

# callback

batch_end_callback = callback.Speedometer(train_data.batch_size, frequent=args.frequent)

means = np.tile(np.array(config.TRAIN.BBOX_MEANS), 2 if config.CLASS_AGNOSTIC else config.dataset.NUM_CLASSES)

stds = np.tile(np.array(config.TRAIN.BBOX_STDS), 2 if config.CLASS_AGNOSTIC else config.dataset.NUM_CLASSES)

epoch_end_callback = [mx.callback.module_checkpoint(mod, prefix, period=1, save_optimizer_states=True), callback.do_checkpoint(prefix, means, stds)]

# decide learning rate

base_lr = lr

lr_factor = config.TRAIN.lr_factor

lr_epoch = [float(epoch) for epoch in lr_step.split(',')]

lr_epoch_diff = [epoch - begin_epoch for epoch in lr_epoch if epoch > begin_epoch]

lr = base_lr * (lr_factor ** (len(lr_epoch) - len(lr_epoch_diff)))

lr_iters = [int(epoch * len(roidb) / batch_size) for epoch in lr_epoch_diff]

print('lr', lr, 'lr_epoch_diff', lr_epoch_diff, 'lr_iters', lr_iters)

lr_scheduler = WarmupMultiFactorScheduler(lr_iters, lr_factor, config.TRAIN.warmup, config.TRAIN.warmup_lr, config.TRAIN.warmup_step)

# optimizer

optimizer_params = {'momentum': config.TRAIN.momentum,

'wd': config.TRAIN.wd,

'learning_rate': lr,

'lr_scheduler': lr_scheduler,

'rescale_grad': 1.0,

'clip_gradient': None}

if not isinstance(train_data, PrefetchingIter):

train_data = PrefetchingIter(train_data)

# train

mod.fit(train_data, eval_metric=eval_metrics, epoch_end_callback=epoch_end_callback,

batch_end_callback=batch_end_callback, kvstore=config.default.kvstore,

optimizer='sgd', optimizer_params=optimizer_params,

print('Called with argument:', args)

ctx = [mx.gpu(int(i)) for i in config.gpus.split(',')]

train_net(args, ctx, config.network.pretrained, config.network.pretrained_epoch, config.TRAIN.model_prefix,