def train_rcnn(cfg,
dataset,
image_set,
root_path,
dataset_path,
frequent,
kvstore,
flip,
shuffle,
resume,
ctx,
pretrained,
epoch,
prefix,
begin_epoch,
end_epoch,
train_shared,
lr,
lr_step,
proposal,
logger=None,
output_path=None):
# set up logger
if not logger:
logging.basicConfig()
logger = logging.getLogger()
logger.setLevel(logging.INFO)
# load symbol
sym_instance = eval(cfg.symbol + '.' + cfg.symbol)()
sym = sym_instance.get_symbol_rfcn(cfg, is_train=True)
# setup multi-gpu
batch_size = len(ctx)
input_batch_size = cfg.TRAIN.BATCH_IMAGES * batch_size
# print cfg
pprint.pprint(cfg)
logger.info('training rcnn cfg:{}\n'.format(pprint.pformat(cfg)))
# load dataset and prepare imdb for training
image_sets = [iset for iset in image_set.split('+')]
roidbs = [
load_proposal_roidb(dataset,
image_set,
root_path,
dataset_path,
proposal=proposal,
append_gt=True,
flip=flip,
result_path=output_path)
for image_set in image_sets
]
roidb = merge_roidb(roidbs)
roidb = filter_roidb(roidb, cfg)
means, stds = add_bbox_regression_targets(roidb, cfg)
# load training data
train_data = ROIIter(roidb,
cfg,
batch_size=input_batch_size,
shuffle=shuffle,
ctx=ctx,
aspect_grouping=cfg.TRAIN.ASPECT_GROUPING)
# infer max shape
max_data_shape = [('data', (cfg.TRAIN.BATCH_IMAGES, 3,
max([v[0] for v in cfg.SCALES]),
max([v[1] for v in cfg.SCALES])))]
# infer shape
data_shape_dict = dict(train_data.provide_data_single +
train_data.provide_label_single)
sym_instance.infer_shape(data_shape_dict)
# load and initialize params
if 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_rfcn(cfg, arg_params, aux_params)
# check parameter shapes
sym_instance.check_parameter_shapes(arg_params, aux_params,
data_shape_dict)
# prepare training
# create solver
data_names = [k[0] for k in train_data.provide_data_single]
label_names = [k[0] for k in train_data.provide_label_single]
if train_shared:
fixed_param_prefix = cfg.network.FIXED_PARAMS_SHARED
else:
fixed_param_prefix = cfg.network.FIXED_PARAMS
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)],
fixed_param_prefix=fixed_param_prefix)
if cfg.TRAIN.RESUME:
mod._preload_opt_states = '%s-%04d.states' % (prefix, begin_epoch)
# decide training params
# metric
eval_metric = metric.RCNNAccMetric(cfg)
cls_metric = metric.RCNNLogLossMetric(cfg)
bbox_metric = metric.RCNNL1LossMetric(cfg)
eval_metrics = mx.metric.CompositeEvalMetric()
for child_metric in [eval_metric, cls_metric, bbox_metric]:
eval_metrics.add(child_metric)
# callback
batch_end_callback = callback.Speedometer(train_data.batch_size,
frequent=frequent)
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 = cfg.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,
cfg.TRAIN.warmup,
cfg.TRAIN.warmup_lr,
cfg.TRAIN.warmup_step)
# optimizer
optimizer_params = {
'momentum': cfg.TRAIN.momentum,
'wd': cfg.TRAIN.wd,
'learning_rate': lr,
'lr_scheduler': lr_scheduler,
'rescale_grad': 1.0,
'clip_gradient': None
}
# train
if not isinstance(train_data, PrefetchingIter):
train_data = PrefetchingIter(train_data)
mod.fit(train_data,
eval_metric=eval_metrics,
epoch_end_callback=epoch_end_callback,
batch_end_callback=batch_end_callback,
kvstore=kvstore,
optimizer='sgd',
optimizer_params=optimizer_params,
arg_params=arg_params,
aux_params=aux_params,
begin_epoch=begin_epoch,
num_epoch=end_epoch)
def train_net(args, ctx, pretrained, pretrained_flow, epoch, prefix,
begin_epoch, end_epoch, lr, lr_step):
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_train_symbol(config)
feat_sym = sym.get_internals()['rpn_cls_score_output']
feat_conv_3x3_relu = sym.get_internals()['feat_conv_3x3_relu_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,
feat_conv_3x3_relu,
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,
normalize_target=config.network.NORMALIZE_RPN,
bbox_mean=config.network.ANCHOR_MEANS,
bbox_std=config.network.ANCHOR_STDS)
# 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]))),
# ('data_ref', (config.TRAIN.BATCH_IMAGES, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES]))),
# ('eq_flag', (1,))]
data_shape1 = {
'data_ref':
(config.TRAIN.BATCH_IMAGES, 3, max([v[0] for v in config.SCALES]),
max([v[1] for v in config.SCALES])),
}
_, feat_shape111, _ = feat_conv_3x3_relu.infer_shape(**data_shape1)
max_data_shape = [('data_ref', (config.TRAIN.BATCH_IMAGES, 3,
max([v[0] for v in config.SCALES]),
max([v[1] for v in config.SCALES]))),
('eq_flag', (1, )),
('motion_vector', (config.TRAIN.BATCH_IMAGES, 2,
int(feat_shape111[0][2]),
int(feat_shape111[0][3])))]
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)
#arg_params_flow, aux_params_flow = load_param(pretrained_flow, epoch, convert=True)
#arg_params.update(arg_params_flow)
#aux_params.update(aux_params_flow)
sym_instance.init_weight(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)
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)
print('Start to train model')
# 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,
arg_params=arg_params,
aux_params=aux_params,
begin_epoch=begin_epoch,
num_epoch=end_epoch)
def train_net(args, ctx, pretrained, pretrained_base, pretrained_ec, epoch,
prefix, begin_epoch, end_epoch, lr, lr_step):
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_train_symbol(config)
# 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('+')]
segdbs = [
load_gt_segdb(config.dataset.dataset,
image_set,
config.dataset.root_path,
config.dataset.dataset_path,
result_path=final_output_path,
flip=config.TRAIN.FLIP) for image_set in image_sets
]
segdb = merge_segdb(segdbs)
# load training data
train_data = TrainDataLoader(sym,
segdb,
config,
batch_size=input_batch_size,
crop_height=config.TRAIN.CROP_HEIGHT,
crop_width=config.TRAIN.CROP_WIDTH,
shuffle=config.TRAIN.SHUFFLE,
ctx=ctx)
# 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]))),
('data_ref', (config.TRAIN.KEY_INTERVAL - 1, 3,
max([v[0] for v in config.SCALES]),
max([v[1] for v in config.SCALES]))),
('eq_flag', (1, ))]
max_data_shape, max_label_shape = train_data.infer_shape(max_data_shape)
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:
print pretrained
arg_params, aux_params = load_param(pretrained, epoch, convert=True)
arg_params_base, aux_params_base = load_param(pretrained_base,
epoch,
convert=True)
arg_params.update(arg_params_base)
aux_params.update(aux_params_base)
arg_params_ec, aux_params_ec = load_param(
pretrained_ec,
epoch,
convert=True,
argprefix=config.TRAIN.arg_prefix)
arg_params.update(arg_params_ec)
aux_params.update(aux_params_ec)
sym_instance.init_weight(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)
if config.TRAIN.RESUME:
mod._preload_opt_states = '%s-%04d.states' % (prefix, begin_epoch)
# decide training params
# metric
fcn_loss_metric = metric.FCNLogLossMetric(config.default.frequent *
batch_size)
eval_metrics = mx.metric.CompositeEvalMetric()
for child_metric in [fcn_loss_metric]:
eval_metrics.add(child_metric)
# callback
batch_end_callback = callback.Speedometer(train_data.batch_size,
frequent=args.frequent)
epoch_end_callback = mx.callback.module_checkpoint(
mod, prefix, period=1, save_optimizer_states=True)
# decide learning rate
base_lr = lr
lr_factor = 0.1
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(segdb) / 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,
arg_params=arg_params,
aux_params=aux_params,
begin_epoch=begin_epoch,
num_epoch=end_epoch)
def train_net(args, ctx, pretrained_dir, pretrained_resnet, epoch, prefix,
begin_epoch, end_epoch, lr, lr_step):
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)))
git_commit_id = commands.getoutput('git rev-parse HEAD')
print("Git commit id:", git_commit_id)
logger.info('Git commit id: {}'.format(git_commit_id))
# 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,
motion_iou_path=config.dataset.motion_iou_path,
flip=config.TRAIN.FLIP,
use_philly=args.usePhilly) 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,
normalize_target=config.network.NORMALIZE_RPN,
bbox_mean=config.network.ANCHOR_MEANS,
bbox_std=config.network.ANCHOR_STDS)
# 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)
# 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)
# load and initialize params
params_loaded = False
if config.TRAIN.RESUME:
arg_params, aux_params = load_param(prefix, begin_epoch, convert=True)
mod._preload_opt_states = '%s-%04d.states' % (prefix, begin_epoch)
print('continue training from ', begin_epoch)
logger.info('continue training from ', begin_epoch)
params_loaded = True
elif config.TRAIN.AUTO_RESUME:
for cur_epoch in range(end_epoch - 1, begin_epoch, -1):
params_filename = '{}-{:04d}.params'.format(prefix, cur_epoch)
states_filename = '{}-{:04d}.states'.format(prefix, cur_epoch)
if os.path.exists(params_filename) and os.path.exists(
states_filename):
begin_epoch = cur_epoch
arg_params, aux_params = load_param(prefix,
cur_epoch,
convert=True)
mod._preload_opt_states = states_filename
print('auto continue training from {}, {}'.format(
params_filename, states_filename))
logger.info('auto continue training from {}, {}'.format(
params_filename, states_filename))
params_loaded = True
break
if not params_loaded:
arg_params, aux_params = load_param(os.path.join(
pretrained_dir, pretrained_resnet),
epoch,
convert=True)
sym_instance.init_weight(config, arg_params, aux_params)
# check parameter shapes
sym_instance.check_parameter_shapes(arg_params, aux_params,
data_shape_dict)
# decide training params
# metric
eval_metric = metric.RCNNAccMetric(config)
cls_metric = metric.RCNNLogLossMetric(config)
bbox_metric = metric.RCNNL1LossMetric(config)
eval_metrics = mx.metric.CompositeEvalMetric()
for child_metric in [eval_metric, cls_metric, bbox_metric]:
eval_metrics.add(child_metric)
if config.TRAIN.JOINT_TRAINING or (not config.TRAIN.LEARN_NMS):
rpn_eval_metric = metric.RPNAccMetric()
rpn_cls_metric = metric.RPNLogLossMetric()
rpn_bbox_metric = metric.RPNL1LossMetric()
for child_metric in [rpn_eval_metric, rpn_cls_metric, rpn_bbox_metric]:
eval_metrics.add(child_metric)
if config.TRAIN.LEARN_NMS:
eval_metrics.add(metric.NMSLossMetric(config, 'pos'))
eval_metrics.add(metric.NMSLossMetric(config, 'neg'))
eval_metrics.add(metric.NMSAccMetric(config))
# callback
batch_end_callback = [
callback.Speedometer(train_data.batch_size, frequent=args.frequent)
]
if config.USE_PHILLY:
total_iter = (end_epoch - begin_epoch) * len(roidb) / input_batch_size
progress_frequent = min(args.frequent * 10, 100)
batch_end_callback.append(
callback.PhillyProgressCallback(total_iter, progress_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 * len(ctx) * config.TRAIN.BATCH_IMAGES
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,
arg_params=arg_params,
aux_params=aux_params,
begin_epoch=begin_epoch,
num_epoch=end_epoch)
def train_rcnn(cfg, dataset, image_set, root_path, dataset_path,
frequent, kvstore, flip, shuffle, resume,
ctx, pretrained, epoch, prefix, begin_epoch, end_epoch,
train_shared, lr, lr_step, proposal, logger=None, output_path=None):
mx.random.seed(3)
np.random.seed(3)
# set up logger
if not logger:
logging.basicConfig()
logger = logging.getLogger()
logger.setLevel(logging.INFO)
# load symbol
sym_instance = eval(cfg.symbol + '.' + cfg.symbol)()
sym = sym_instance.get_symbol_rcnn(cfg, is_train=True)
# setup multi-gpu
batch_size = len(ctx)
input_batch_size = cfg.TRAIN.BATCH_IMAGES * batch_size
# print cfg
pprint.pprint(cfg)
logger.info('training rcnn cfg:{}\n'.format(pprint.pformat(cfg)))
# load dataset and prepare imdb for training
image_sets = [iset for iset in image_set.split('+')]
roidbs = [load_proposal_roidb(dataset, image_set, root_path, dataset_path,
proposal=proposal, append_gt=True, flip=flip, result_path=output_path)
for image_set in image_sets]
roidb = merge_roidb(roidbs)
roidb = filter_roidb(roidb, cfg)
means, stds = add_bbox_regression_targets(roidb, cfg)
# load training data
train_data = ROIIter(roidb, cfg, batch_size=input_batch_size, shuffle=shuffle,
ctx=ctx, aspect_grouping=cfg.TRAIN.ASPECT_GROUPING)
# infer max shape
max_data_shape = [('data', (cfg.TRAIN.BATCH_IMAGES, 3, max([v[0] for v in cfg.SCALES]), max([v[1] for v in cfg.SCALES])))]
# infer shape
data_shape_dict = dict(train_data.provide_data_single + train_data.provide_label_single)
sym_instance.infer_shape(data_shape_dict)
# load and initialize params
if 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_rcnn(cfg, arg_params, aux_params)
# check parameter shapes
sym_instance.check_parameter_shapes(arg_params, aux_params, data_shape_dict)
# prepare training
# create solver
data_names = [k[0] for k in train_data.provide_data_single]
label_names = [k[0] for k in train_data.provide_label_single]
if train_shared:
fixed_param_prefix = cfg.network.FIXED_PARAMS_SHARED
else:
fixed_param_prefix = cfg.network.FIXED_PARAMS
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)], fixed_param_prefix=fixed_param_prefix)
if cfg.TRAIN.RESUME:
mod._preload_opt_states = '%s-%04d.states'%(prefix, begin_epoch)
# decide training params
# metric
eval_metric = metric.RCNNAccMetric(cfg)
cls_metric = metric.RCNNLogLossMetric(cfg)
bbox_metric = metric.RCNNL1LossMetric(cfg)
eval_metrics = mx.metric.CompositeEvalMetric()
for child_metric in [eval_metric, cls_metric, bbox_metric]:
eval_metrics.add(child_metric)
# callback
batch_end_callback = callback.Speedometer(train_data.batch_size, frequent=frequent)
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 = cfg.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, cfg.TRAIN.warmup, cfg.TRAIN.warmup_lr, cfg.TRAIN.warmup_step)
# optimizer
optimizer_params = {'momentum': cfg.TRAIN.momentum,
'wd': cfg.TRAIN.wd,
'learning_rate': lr,
'lr_scheduler': lr_scheduler,
'rescale_grad': 1.0,
'clip_gradient': None}
# train
if not isinstance(train_data, PrefetchingIter):
train_data = PrefetchingIter(train_data)
mod.fit(train_data, eval_metric=eval_metrics, epoch_end_callback=epoch_end_callback,
batch_end_callback=batch_end_callback, kvstore=kvstore,
optimizer='sgd', optimizer_params=optimizer_params,
arg_params=arg_params, aux_params=aux_params, begin_epoch=begin_epoch, num_epoch=end_epoch)
def train_net(args, ctx, pretrained, epoch, prefix, begin_epoch, end_epoch, lr, lr_step):
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)
# 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)
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,
arg_params=arg_params, aux_params=aux_params, begin_epoch=begin_epoch, num_epoch=end_epoch)
def train_net(args, ctx, pretrained, epoch, prefix, begin_epoch, end_epoch, lr,
lr_step):
# 创建logger和对应的输出路径
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)
# 特征symbol，从网络sym中获取rpn_cls_score_output
feat_sym = sym.get_internals()['rpn_cls_score_output']
# setup multi-gpu
# 使能多GPU训练，每一张卡训练一个batch
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
# 加载数据集同时准备训练的imdb，使用+分割不同的图像数据集，比如2007_trainval+2012_trainval
image_sets = [iset for iset in config.dataset.image_set.split('+')]
# load gt roidb加载gt roidb，根据数据集类型，图像集具体子类，数据集根目录和数据集路径，同时配置相关TRAIN为FLIP来增广数据
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
roidb = merge_roidb(roidbs)
# 根据配置文件中对应的过滤规则来滤出roi
roidb = filter_roidb(roidb, config)
# load training data
# 加载训练数据，anchor Loader为对应分类和回归的锚点加载，通过对应的roidb，查找对应的正负样本的锚点，该生成器需要参数锚点尺度，ratios和对应的feature的stride
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
# 加载并且初始化参数，如果训练中是继续上次的训练，也就是RESUME这一flag设置为True
if config.TRAIN.RESUME:
print('continue training from ', begin_epoch)
# 从前缀和being_epoch中加载RESUME的arg参数和aux参数，同时需要转换为GPU NDArray
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)
# check parameter shapes
# 检查相关参数的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)
if config.TRAIN.RESUME:
mod._preload_opt_states = '%s-%04d.states' % (prefix, begin_epoch)
# decide training params
# metric
# 以下主要是RPN和RCNN相关的一些评价指标
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)
# mxnet中合成的评估指标，可以增加以上所有的评估指标，包括rpn_eval_metrix、rpn_cls_metric、rpn_bbox_metric和rcnn_eval_metric、rcnn_cls_metric、rcnn_bbox_metric
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后的callback回调以及epoch后的callback回调
# batch_end_callback是在训练一定batch_size后进行的相应回调，回调频率为frequent
batch_end_callback = callback.Speedometer(train_data.batch_size,
frequent=args.frequent)
# means和stds，如果BBOX是类无关的，那么means为复制means两个，否则复制数量为NUM_CLASSES
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为一个周期结束后的回调
epoch_end_callback = [
mx.callback.module_checkpoint(mod,
prefix,
period=1,
save_optimizer_states=True),
callback.do_checkpoint(prefix, means, stds)
]
# decide learning rate
# 以下主要根据不同的学习率调整策略来决定学习率，这里如voc中默认的初始lr为0.0005
base_lr = lr
# 学习率调整因子
lr_factor = config.TRAIN.lr_factor
# 学习率调整周期，lr_step一般格式为3, 5，表示在3和5周期中进行学习率调整
lr_epoch = [float(epoch) for epoch in lr_step.split(',')]
# 如果当前周期大于begin_epoch那么lr_epoch_diff为epoch-begin_epoch
lr_epoch_diff = [
epoch - begin_epoch for epoch in lr_epoch if epoch > begin_epoch
]
print('lr_epoch', lr_epoch, 'begin_epoch', begin_epoch)
# 通过当前的epoch来计算当前应该具有的lr
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)
# learning rate调整机制，warmup multi factor scheduler预训练多因子调整器
lr_scheduler = WarmupMultiFactorScheduler(lr_iters, lr_factor,
config.TRAIN.warmup,
config.TRAIN.warmup_lr,
config.TRAIN.warmup_step)
# optimizer
# 优化器参数，包含momentum、wd、lr、lr_scheduler、rescale_grad和clip_gradient
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):
print('!!!train_data is not PrefetchingIter!!!')
train_data = PrefetchingIter(train_data)
# train
# 模型训练过程，输入train_data，评估指标包括eval_metrics等一系列指标，每一个epoch结束后进入epoch_end_callback，每一个batch结束后进入batch_end_callback，优化器使用sgd，同时优化参数、输入参数和辅助参数以及begin周期和end周期
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,
arg_params=arg_params,
aux_params=aux_params,
begin_epoch=begin_epoch,
num_epoch=end_epoch)
def train_net(args, ctx, pretrained, epoch, prefix, begin_epoch, end_epoch, lr, lr_step):
np.random.seed(0)
mx.random.seed(0)
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_dats_shape=['data', (1,3,600,1000)]
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_lable_shape=[]
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)))
logger.info('providing maximum shape'+str(max_data_shape)+" "+str(max_label_shape))
data_shape_dict = dict(train_data.provide_data_single + train_data.provide_label_single)
# add by chaojie
logger.info("data_sahpe_dict:\n{}".format(pprint.pformat(data_shape_dict)))
pprint.pprint(data_shape_dict)
sym_instance.infer_shape(data_shape_dict)
pprint.pprint(sym_instance.arg_shape_dict)
logger.info("sym_instance.arg_shape_dict\n")
logging.info(pprint.pformat(sym_instance.arg_shape_dict))
#dot = mx.viz.plot_network(sym, node_attrs={'shape': 'rect', 'fixedsize': 'false'})
#dot.render(os.path.join('./output/rcnn/network_vis', config.symbol + '_rcnn'))
# 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)
# 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)
if config.TRAIN.RESUME:
mod._preload_opt_states = '%s-%04d.states'%(prefix, begin_epoch)
# decide training params
# metric
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
if config.TRAIN.JOINT_TRAINING or (not config.TRAIN.LEARN_NMS):
rpn_eval_metric = metric.RPNAccMetric()
rpn_cls_metric = metric.RPNLogLossMetric()
rpn_bbox_metric = metric.RPNL1LossMetric()
for child_metric in [rpn_eval_metric, rpn_cls_metric, rpn_bbox_metric]:
eval_metrics.add(child_metric)
for child_metric in [eval_metric, cls_metric, bbox_metric]:
eval_metrics.add(child_metric)
if config.TRAIN.LEARN_NMS:
eval_metrics.add(metric.NMSLossMetric(config, 'pos'))
eval_metrics.add(metric.NMSLossMetric(config, 'neg'))
eval_metrics.add(metric.NMSAccMetric(config))
# 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,
arg_params=arg_params, aux_params=aux_params, begin_epoch=begin_epoch, num_epoch=end_epoch)
def train_net(args, ctx, pretrained, epoch, prefix, begin_epoch, end_epoch, lr,
lr_step):
mx.random.seed(3)
np.random.seed(3)
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_pyramid_level = np.log2(config.network.RPN_FEAT_STRIDE).astype(int)
feat_sym = [
sym.get_internals()['rpn_cls_score_p' + str(x) + '_output']
for x in feat_pyramid_level
]
print('load symbol END')
# 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
print('Start load dataset and prepare imdb for training')
image_sets = [iset for iset in config.dataset.image_set.split('+')]
roidbs = [
load_gt_roidb_poly(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)
print('Start load training data')
# load training data
train_data = PyramidAnchorIterator_poly(
feat_sym,
roidb,
config,
batch_size=input_batch_size,
shuffle=config.TRAIN.SHUFFLE,
ctx=ctx,
feat_strides=config.network.RPN_FEAT_STRIDE,
anchor_scales=config.network.ANCHOR_SCALES,
anchor_ratios=config.network.ANCHOR_RATIOS,
aspect_grouping=config.TRAIN.ASPECT_GROUPING,
allowed_border=np.inf)
# 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, 300, 9)))
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)
# 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)
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()
rpn_fg_metric = metric.RPNFGFraction(config)
eval_fg_metric = metric.RCNNFGAccuracy(config)
eval_metric = metric.RCNNAccMetric(config)
cls_metric = metric.RCNNLogLossMetric(config)
bbox_metric = metric.RCNNL1LossMetric(config)
# add Rroi loss here
RCNN_proposal_fraction_metric = metric.RCNNFGFraction(config)
Rroi_fg_accuracy = metric.RRoIRCNNFGAccuracy(config)
Rroi_accuracy = metric.RRoIAccMetric(config)
Rroi_cls_metric = metric.RRoIRCNNLogLossMetric(config)
Rroi_bbox_metric = metric.RRoIRCNNL1LossMetric(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, rpn_fg_metric,
eval_fg_metric, eval_metric, cls_metric, bbox_metric,
RCNN_proposal_fraction_metric, Rroi_fg_accuracy, Rroi_accuracy,
Rroi_cls_metric, Rroi_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)
Rroi_means = np.tile(
np.array(config.TRAIN.RRoI_BBOX_MEANS), 2
if config.network.RRoI_CLASS_AGNOSTIC else config.dataset.NUM_CLASSES)
Rroi_stds = np.tile(
np.array(config.TRAIN.RRoI_BBOX_STDS), 2
if config.network.RRoI_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_Rroi(prefix, means, stds, Rroi_means, Rroi_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,
'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,
arg_params=arg_params,
aux_params=aux_params,
begin_epoch=begin_epoch,
num_epoch=end_epoch)
def train_net(args, ctx, pretrained, epoch, prefix, begin_epoch, end_epoch, lr,
lr_step):
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,
arg_params=arg_params,
aux_params=aux_params,
begin_epoch=begin_epoch,
num_epoch=end_epoch)
def train_net(args, ctx, pretrained, epoch, prefix, begin_epoch, end_epoch, lr, lr_step):
# 创建logger和对应的输出路径
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)
# 特征symbol，从网络sym中获取rpn_cls_score_output
feat_sym = sym.get_internals()['rpn_cls_score_output']
# setup multi-gpu
# 使能多GPU训练，每一张卡训练一个batch
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
# 加载数据集同时准备训练的imdb，使用+分割不同的图像数据集，比如2007_trainval+2012_trainval
image_sets = [iset for iset in config.dataset.image_set.split('+')]
# load gt roidb加载gt roidb，根据数据集类型，图像集具体子类，数据集根目录和数据集路径，同时配置相关TRAIN为FLIP来增广数据
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
roidb = merge_roidb(roidbs)
# 根据配置文件中对应的过滤规则来滤出roi
roidb = filter_roidb(roidb, config)
# load training data
# 加载训练数据，anchor Loader为对应分类和回归的锚点加载，通过对应的roidb，查找对应的正负样本的锚点，该生成器需要参数锚点尺度，ratios和对应的feature的stride
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)
# 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)
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,
arg_params=arg_params, aux_params=aux_params, begin_epoch=begin_epoch, num_epoch=end_epoch)
def train_rcnn(cfg, dataset, image_set, root_path, dataset_path,
frequent, kvstore, flip, shuffle, resume,
ctx, pretrained, epoch, prefix, begin_epoch, end_epoch,
train_shared, lr, lr_step, proposal, logger=None, output_path=None):
mx.random.seed(0)
np.random.seed(0)
# set up logger
if not logger:
logging.basicConfig()
logger = logging.getLogger()
logger.setLevel(logging.INFO)
# load symbol
sym_instance = eval(cfg.symbol + '.' + cfg.symbol)()
sym = sym_instance.get_symbol_rcnn(cfg, is_train=True)
# setup multi-gpu
batch_size = len(ctx)
input_batch_size = cfg.TRAIN.BATCH_IMAGES * batch_size
# print cfg
pprint.pprint(cfg)
logger.info('training rcnn cfg:{}\n'.format(pprint.pformat(cfg)))
rpn_path = cfg.dataset.proposal_cache
# load dataset and prepare imdb for training
image_sets = [iset for iset in image_set.split('+')]
roidbs = [load_proposal_roidb(dataset, image_set, root_path, dataset_path,
proposal=proposal, append_gt=True, flip=flip, result_path=output_path,
rpn_path=rpn_path, top_roi=cfg.TRAIN.TOP_ROIS)
for image_set in image_sets]
roidb = merge_roidb(roidbs)
roidb = filter_roidb(roidb, cfg)
means, stds = add_bbox_regression_targets(roidb, cfg)
# load training data
train_data = ROIIter(roidb, cfg, batch_size=input_batch_size, shuffle=shuffle,
ctx=ctx, aspect_grouping=cfg.TRAIN.ASPECT_GROUPING)
# infer max shape
max_height = max([v[0] for v in cfg.SCALES])
max_width = max([v[1] for v in cfg.SCALES])
paded_max_height = max_height + cfg.network.IMAGE_STRIDE - max_height % cfg.network.IMAGE_STRIDE
paded_max_width = max_width + cfg.network.IMAGE_STRIDE - max_width % (cfg.network.IMAGE_STRIDE)
max_data_shape = [('data', (cfg.TRAIN.BATCH_IMAGES, 3, paded_max_height, paded_max_width))]
# infer shape
data_shape_dict = dict(train_data.provide_data_single + train_data.provide_label_single)
sym_instance.infer_shape(data_shape_dict)
# print shape
pprint.pprint(sym_instance.arg_shape_dict)
logging.info(pprint.pformat(sym_instance.arg_shape_dict))
max_batch_roi = cfg.TRAIN.TOP_ROIS if cfg.TRAIN.BATCH_ROIS == -1 else cfg.TRAIN.BATCH_ROIS
num_class = 2 if cfg.CLASS_AGNOSTIC else cfg.dataset.NUM_CLASSES
max_label_shape = [('label', (cfg.TRAIN.BATCH_IMAGES, max_batch_roi)),
('bbox_target', (cfg.TRAIN.BATCH_IMAGES, max_batch_roi, num_class * 4)),
('bbox_weight', (cfg.TRAIN.BATCH_IMAGES, max_batch_roi, num_class * 4))]
if cfg.network.USE_NONGT_INDEX:
max_label_shape.append(('nongt_index', (2000,)))
if cfg.network.ROIDispatch:
max_data_shape.append(('rois_0', (cfg.TRAIN.BATCH_IMAGES, max_batch_roi / 4, 5)))
max_data_shape.append(('rois_1', (cfg.TRAIN.BATCH_IMAGES, max_batch_roi / 4, 5)))
max_data_shape.append(('rois_2', (cfg.TRAIN.BATCH_IMAGES, max_batch_roi / 4, 5)))
max_data_shape.append(('rois_3', (cfg.TRAIN.BATCH_IMAGES, max_batch_roi / 4, 5)))
else:
max_data_shape.append(('rois', (cfg.TEST.PROPOSAL_POST_NMS_TOP_N + 30, 5)))
#dot = mx.viz.plot_network(sym, node_attrs={'shape': 'rect', 'fixedsize': 'false'})
#dot.render(os.path.join('./output/rcnn/network_vis', cfg.symbol + cfg.TRAIN.model_prefix))
# load and initialize params
if 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_rcnn(cfg, arg_params, aux_params)
# check parameter shapes
sym_instance.check_parameter_shapes(arg_params, aux_params, data_shape_dict)
# prepare training
# create solver
data_names = [k[0] for k in train_data.provide_data_single]
label_names = [k[0] for k in train_data.provide_label_single]
if train_shared:
fixed_param_prefix = cfg.network.FIXED_PARAMS_SHARED
else:
fixed_param_prefix = cfg.network.FIXED_PARAMS
if cfg.network.ROIDispatch:
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)
else:
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)
if cfg.TRAIN.RESUME:
mod._preload_opt_states = '%s-%04d.states' % (prefix, begin_epoch)
# decide training params
# metric
eval_metric = metric.RCNNAccMetric(cfg)
cls_metric = metric.RCNNLogLossMetric(cfg)
bbox_metric = metric.RCNNL1LossMetric(cfg)
eval_metrics = mx.metric.CompositeEvalMetric()
for child_metric in [eval_metric, cls_metric, bbox_metric]:
eval_metrics.add(child_metric)
if cfg.TRAIN.LEARN_NMS:
eval_metrics.add(metric.NMSLossMetric(cfg, 'pos'))
eval_metrics.add(metric.NMSLossMetric(cfg, 'neg'))
eval_metrics.add(metric.NMSAccMetric(cfg))
# callback
batch_end_callback = callback.Speedometer(train_data.batch_size, frequent=frequent)
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 = cfg.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, cfg.TRAIN.warmup, cfg.TRAIN.warmup_lr,
cfg.TRAIN.warmup_step)
# optimizer
optimizer_params = {'momentum': cfg.TRAIN.momentum,
'wd': cfg.TRAIN.wd,
'learning_rate': lr,
'lr_scheduler': lr_scheduler,
'rescale_grad': 1.0,
'clip_gradient': None}
# train
if not isinstance(train_data, PrefetchingIter):
train_data = PrefetchingIter(train_data)
mod.fit(train_data, eval_metric=eval_metrics, epoch_end_callback=epoch_end_callback,
batch_end_callback=batch_end_callback, kvstore=kvstore,
optimizer='sgd', optimizer_params=optimizer_params,
arg_params=arg_params, aux_params=aux_params, begin_epoch=begin_epoch, num_epoch=end_epoch)
def train_net(args, ctx, pretrained, epoch, prefix, begin_epoch, end_epoch, lr,
lr_step):
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_retina_symbol(config, is_train=True)
feat_sym = []
feat_sym_p4 = sym.get_internals()['box_pred/p4_output']
feat_sym_p5 = sym.get_internals()['box_pred/p5_output']
feat_sym_p6 = sym.get_internals()['box_pred/p6_output']
feat_sym_p7 = sym.get_internals()['box_pred/p7_output']
feat_sym.append(feat_sym_p4)
feat_sym.append(feat_sym_p5)
feat_sym.append(feat_sym_p6)
feat_sym.append(feat_sym_p7)
#######
feat_stride = []
feat_stride.append(config.network.p4_RPN_FEAT_STRIDE)
feat_stride.append(config.network.p5_RPN_FEAT_STRIDE)
feat_stride.append(config.network.p6_RPN_FEAT_STRIDE)
feat_stride.append(config.network.p7_RPN_FEAT_STRIDE)
anchor_scales = []
anchor_scales.append(config.network.p4_ANCHOR_SCALES)
anchor_scales.append(config.network.p5_ANCHOR_SCALES)
anchor_scales.append(config.network.p6_ANCHOR_SCALES)
anchor_scales.append(config.network.p7_ANCHOR_SCALES)
anchor_ratios = []
anchor_ratios.append(config.network.p4_ANCHOR_RATIOS)
anchor_ratios.append(config.network.p5_ANCHOR_RATIOS)
anchor_ratios.append(config.network.p6_ANCHOR_RATIOS)
anchor_ratios.append(config.network.p7_ANCHOR_RATIOS)
#############
# 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,
feat_stride,
anchor_scales,
anchor_ratios,
roidb,
config,
batch_size=input_batch_size,
shuffle=config.TRAIN.SHUFFLE,
ctx=ctx,
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
# infer max 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)
# 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)
if config.TRAIN.RESUME:
mod._preload_opt_states = '%s-%04d.states' % (prefix, begin_epoch)
# decide training params
# metric
Retina_toal_eval_metric = metric.RetinaToalAccMetric()
Retina_cls_metric = metric.RetinaFocalLossMetric()
Retina_bbox_metric = metric.RetinaL1LossMetric()
eval_metrics = mx.metric.CompositeEvalMetric()
# rpn_eval_metric, rpn_cls_metric, rpn_bbox_metric, eval_metric, cls_metric, bbox_metric
for child_metric in [
Retina_toal_eval_metric, Retina_cls_metric, Retina_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_step.split(',')
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 = {
'learning_rate': lr,
'wd': 0.0001,
}
if not isinstance(train_data, PrefetchingIter):
train_data = PrefetchingIter(train_data)
# train
initializer = mx.init.MSRAPrelu(factor_type='out', slope=0)
# adam = mx.optimizer.AdaDelta(rho=0.09, epsilon=1e-14)
#optimizer_params=optimizer_params,
print "-----------------------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='adam',
optimizer_params=optimizer_params,
initializer=initializer,
arg_params=arg_params,
aux_params=aux_params,
begin_epoch=begin_epoch,
num_epoch=end_epoch)
def train_net(args, ctx, pretrained, epoch, prefix, begin_epoch, end_epoch, lr,
lr_step):
if config.dataset.dataset != 'JSONList':
logger, final_output_path = create_logger(config.output_path, args.cfg,
config.dataset.image_set)
prefix = os.path.join(final_output_path, prefix)
else:
import datetime
import logging
final_output_path = config.output_path
prefix = prefix + '_' + datetime.datetime.now().strftime(
"%Y-%m-%d_%H_%M_%S")
prefix = os.path.join(final_output_path, prefix)
shutil.copy2(args.cfg, prefix + '.yaml')
log_file = prefix + '.log'
head = '%(asctime)-15s %(message)s'
logging.basicConfig(filename=log_file, format=head)
logger = logging.getLogger()
logger.setLevel(logging.INFO)
logger.info('prefix: %s' % prefix)
print('prefix: %s' % 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)
# 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
if config.network.MULTI_RPN:
num_layers = len(config.network.MULTI_RPN_STRIDES)
rpn_syms = [
sym.get_internals()['rpn%d_cls_score_output' % l]
for l in range(num_layers)
]
train_data = PyramidAnchorLoader(
rpn_syms,
roidb,
config,
batch_size=input_batch_size,
shuffle=config.TRAIN.SHUFFLE,
ctx=ctx,
feat_strides=config.network.MULTI_RPN_STRIDES,
anchor_scales=config.network.ANCHOR_SCALES,
anchor_ratios=config.network.ANCHOR_RATIOS,
aspect_grouping=config.TRAIN.ASPECT_GROUPING,
allowed_border=np.inf)
else:
feat_sym = sym.get_internals()['rpn_cls_score_output']
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)
# 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)
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)
if config.network.PREDICT_KEYPOINTS:
kps_cls_acc = metric.KeypointAccMetric(config)
kps_cls_loss = metric.KeypointLogLossMetric(config)
kps_pos_loss = metric.KeypointL1LossMetric(config)
eval_metrics.add(kps_cls_acc)
eval_metrics.add(kps_cls_loss)
eval_metrics.add(kps_pos_loss)
# 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,
arg_params=arg_params,
aux_params=aux_params,
begin_epoch=begin_epoch,
num_epoch=end_epoch)
def train_feature_distance_net(args, ctx, pretrained, pretrained_flow, epoch,
prefix, begin_epoch, end_epoch, lr, lr_step):
# ==============prepare logger==============
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)()
if config.TRAIN.G_type == 0:
sym = sym_instance.get_train_feature_distance_symbol(config)
elif config.TRAIN.G_type == 1:
sym = sym_instance.get_train_feature_distance_symbol_res(config)
# ==============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)
train_iter = ImagenetVIDIter(roidb, input_batch_size, config,
config.TRAIN.SHUFFLE, ctx)
# 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]))),
('data_ref', (config.TRAIN.BATCH_IMAGES, 3,
max([v[0] for v in config.SCALES]),
max([v[1] for v in config.SCALES])))]
print 'providing maximum shape', max_data_shape
data_shape_dict = dict(train_iter.provide_data_single)
pprint.pprint(data_shape_dict)
sym_instance.infer_shape(data_shape_dict)
# ==============init 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)
# arg_params_flow, aux_params_flow = load_param(pretrained_flow, epoch, convert=True)
# arg_params.update(arg_params_flow)
# aux_params.update(aux_params_flow)
sym_instance.init_weight(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 = train_iter.data_name
label_names = train_iter.label_name
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=None,
fixed_param_prefix=fixed_param_prefix)
if config.TRAIN.RESUME:
mod._preload_opt_states = '%s-%04d.states' % (prefix, begin_epoch)
# ==============optimizer==============
optimizer_params = {
'learning_rate': 0.00005,
}
if not isinstance(train_iter, PrefetchingIter):
train_iter = PrefetchingIter(train_iter)
batch_end_callback = callback.Speedometer(train_iter.batch_size,
frequent=args.frequent)
epoch_end_callback = [
mx.callback.module_checkpoint(mod,
prefix,
period=1,
save_optimizer_states=True),
callback.do_checkpoint(prefix)
]
feature_L2_loss = metric.FeatureL2LossMetric(config)
eval_metrics = mx.metric.CompositeEvalMetric()
eval_metrics.add(feature_L2_loss)
mod.fit(train_iter,
eval_metric=eval_metrics,
epoch_end_callback=epoch_end_callback,
batch_end_callback=batch_end_callback,
kvstore=config.default.kvstore,
optimizer='RMSprop',
optimizer_params=optimizer_params,
arg_params=arg_params,
aux_params=aux_params,
begin_epoch=begin_epoch,
num_epoch=end_epoch,
initializer=mx.init.Normal(0.02),
allow_missing=True)