def __init__(self,
module,
imdb,
roidb,
test_iter,
cfg,
rcnn_output_names=None,
rpn_output_names=None,
logger=None,
batch_size=None):
self.test_iter = test_iter
# Make sure that iterator is instnace of Prefetching iterator
if test_iter and not isinstance(test_iter, PrefetchingIter):
self.test_iter = PrefetchingIter(self.test_iter)
self.scale = test_iter.test_scale
self.cfg = cfg
self.module = module
if test_iter:
self.data_names = [k[0] for k in test_iter.provide_data_single]
self.rcnn_output_names = rcnn_output_names
if not self.rcnn_output_names:
self.rcnn_output_names = {
'cls': 'cls_prob_reshape_output',
'bbox': 'bbox_pred_reshape_output',
'im_ids': 'im_ids',
'scale_map': 'scale_prob_output',
'im_info': 'im_info',
'chip_ids': 'chip_ids'
}
self.rpn_output_names = rpn_output_names
if not self.rpn_output_names:
self.rpn_output_names = {
'scores': 'rois_score',
'rois': 'rois_output',
'im_ids': 'im_ids'
}
self.logger = logger
self.result_path = imdb.result_path
self.num_classes = imdb.num_classes
self.class_names = imdb.classes
self.num_images = len(roidb)
self.imdb_name = imdb.name
self.nms_worker = nms_worker(cfg.TEST.NMS, cfg.TEST.NMS_SIGMA)
self.batch_size = batch_size
self.roidb = roidb
self.verbose = len(roidb) > 1
self.thread_pool = None
if not self.batch_size:
self.batch_size = self.cfg.TEST.BATCH_IMAGES
bbox_metric = metric.RPNL1LossMetric()
rceval_metric = metric.RCNNAccMetric(config)
rccls_metric = metric.RCNNLogLossMetric(config)
rcbbox_metric = metric.RCNNL1LossCRCNNMetric(config)
eval_metrics = mx.metric.CompositeEvalMetric()
eval_metrics.add(eval_metric)
eval_metrics.add(cls_metric)
eval_metrics.add(bbox_metric)
eval_metrics.add(rceval_metric)
eval_metrics.add(rccls_metric)
eval_metrics.add(rcbbox_metric)
if config.TRAIN.WITH_MASK:
mask_metric = metric.MaskLogLossMetric(config)
eval_metrics.add(mask_metric)
optimizer_params = get_optim_params(config, len(train_iter), batch_size)
print ('Optimizer params: {}'.format(optimizer_params))
# Checkpointing
prefix = os.path.join(output_path, args.save_prefix)
batch_end_callback = mx.callback.Speedometer(batch_size, args.display)
epoch_end_callback = [mx.callback.module_checkpoint(mod, prefix, period=1, save_optimizer_states=True),
eval('{}.checkpoint_callback'.format(config.symbol))(sym_inst.get_bbox_param_names(), prefix, bbox_means, bbox_stds)]
train_iter = PrefetchingIter(train_iter)
mod.fit(train_iter, optimizer='sgd', optimizer_params=optimizer_params,
eval_metric=eval_metrics, num_epoch=config.TRAIN.end_epoch, kvstore=config.default.kvstore,
batch_end_callback=batch_end_callback,
epoch_end_callback=epoch_end_callback, arg_params=arg_params, aux_params=aux_params)