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, tb_log_path = create_logger(config.output_path, config.log_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']
# 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,
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_bef', (config.TRAIN.BATCH_IMAGES, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES]))),
('data_aft', (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)
begin_batch = 0
# load and initialize params
if config.TRAIN.RESUME:
begin_epoch, begin_batch = get_begin_point(final_output_path)
print('continue training from ', begin_epoch, begin_batch)
arg_params, aux_params = load_param(prefix, begin_batch, 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_batch)
# 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
tensorboard_callback = callback.TensorboardCallback(tb_log_path) if config.use_tensorboard else None
batch_end_callback = [callback.Speedometer(train_data.batch_size, frequent=args.frequent)]
if tensorboard_callback:
batch_end_callback.append(tensorboard_callback)
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,
'begin_num_update': begin_batch,
'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, begin_batch = begin_batch)
def train_net(args, ctx, pretrained, epoch, prefix, begin_epoch, end_epoch, lr,
lr_step):
logger, final_output_path, _, tensorboard_path = create_env(
config.output_path, args.cfg, config.dataset.image_set)
prefix = os.path.join(final_output_path, prefix)
# print config
pprint.pprint(config)
logger.info('training config:{}\n'.format(pprint.pformat(config)))
print "config.symbol", config.symbol
sym_instance = eval(config.symbol + '.' + config.symbol)()
if config.TRAIN.use_dynamic:
sym_gen = sym_instance.sym_gen(config, is_train=True)
else:
sym = sym_instance.get_symbol(config, is_train=True)
# infer max shape
scales = [(config.TRAIN.crop_size[0], config.TRAIN.crop_size[1])
] if config.TRAIN.enable_crop else config.SCALES
label_stride = config.network.LABEL_STRIDE
network_ratio = config.network.ratio
if config.network.use_context:
if config.network.use_crop_context:
max_data_shape = [
('data',
(config.TRAIN.BATCH_IMAGES, 3, config.TRAIN.crop_size[0],
config.TRAIN.crop_size[1])),
('origin_data', (config.TRAIN.BATCH_IMAGES, 3, 736, 736)),
('rois', (config.TRAIN.BATCH_IMAGES, 5))
]
else:
max_data_shape = [
('data',
(config.TRAIN.BATCH_IMAGES, 3, config.TRAIN.crop_size[0],
config.TRAIN.crop_size[1])),
('origin_data', (config.TRAIN.BATCH_IMAGES, 3,
int(config.SCALES[0][0] * network_ratio),
int(config.SCALES[0][1] * network_ratio))),
('rois', (config.TRAIN.BATCH_IMAGES, 5))
]
else:
if config.TRAIN.enable_crop:
max_data_shape = [('data', (config.TRAIN.BATCH_IMAGES, 3,
config.TRAIN.crop_size[0],
config.TRAIN.crop_size[1]))]
else:
max_data_shape = [
('data',
(config.TRAIN.BATCH_IMAGES, 3,
max([make_divisible(v[0], label_stride) for v in scales]),
max([make_divisible(v[1], label_stride) for v in scales])))
]
if config.network.use_mult_label:
if config.network.use_crop_context:
max_label_shape = [
('label',
(config.TRAIN.BATCH_IMAGES, 1,
make_divisible(config.TRAIN.crop_size[0], label_stride) //
config.network.LABEL_STRIDE,
make_divisible(config.TRAIN.crop_size[1], label_stride) //
config.network.LABEL_STRIDE)),
('origin_label', (config.TRAIN.BATCH_IMAGES, 1, 736, 736))
]
else:
max_label_shape = [
('label',
(config.TRAIN.BATCH_IMAGES, 1,
make_divisible(config.TRAIN.crop_size[0], label_stride) //
config.network.LABEL_STRIDE,
make_divisible(config.TRAIN.crop_size[1], label_stride) //
config.network.LABEL_STRIDE)),
('origin_label', (config.TRAIN.BATCH_IMAGES, 1,
int(config.SCALES[0][0] * network_ratio),
int(config.SCALES[0][1] * network_ratio)))
]
elif config.network.use_metric:
scale_list = config.network.scale_list
scale_name = ['a', 'b', 'c']
if config.network.scale_list == [1, 2, 4]:
scale_name = ['', '', '']
if config.TRAIN.enable_crop:
if config.TRAIN.use_mult_metric:
max_label_shape = [
('label', (config.TRAIN.BATCH_IMAGES, 1,
config.TRAIN.crop_size[0] // label_stride,
config.TRAIN.crop_size[1] // label_stride)),
('metric_label_' + str(scale_list[0]) + scale_name[0],
(config.TRAIN.BATCH_IMAGES, 9, 1,
config.TRAIN.crop_size[0] // label_stride,
config.TRAIN.crop_size[1] // label_stride)),
('metric_label_' + str(scale_list[1]) + scale_name[1],
(config.TRAIN.BATCH_IMAGES, 9, 1,
config.TRAIN.crop_size[0] // label_stride,
config.TRAIN.crop_size[1] // label_stride)),
('metric_label_' + str(scale_list[2]) + scale_name[2],
(config.TRAIN.BATCH_IMAGES, 9, 1,
config.TRAIN.crop_size[0] // label_stride,
config.TRAIN.crop_size[1] // label_stride))
]
else:
max_label_shape = [
('label', (config.TRAIN.BATCH_IMAGES, 1,
config.TRAIN.crop_size[0] // label_stride,
config.TRAIN.crop_size[1] // label_stride)),
('metric_label',
(config.TRAIN.BATCH_IMAGES, 9, 1,
config.TRAIN.crop_size[0] // label_stride,
config.TRAIN.crop_size[1] // label_stride))
]
else:
if config.TRAIN.use_mult_metric:
max_label_shape = [
('label',
(config.TRAIN.BATCH_IMAGES, 1,
max([make_divisible(v[0], label_stride)
for v in scales]) // config.network.LABEL_STRIDE,
max([make_divisible(v[1], label_stride)
for v in scales]) // config.network.LABEL_STRIDE)),
('metric_label_' + str(scale_list[0]) + scale_name[0],
(config.TRAIN.BATCH_IMAGES, 9, 1,
max([make_divisible(v[0], label_stride)
for v in scales]) // config.network.LABEL_STRIDE,
max([make_divisible(v[1], label_stride)
for v in scales]) // config.network.LABEL_STRIDE)),
('metric_label_' + str(scale_list[1]) + scale_name[1],
(config.TRAIN.BATCH_IMAGES, 9, 1,
max([make_divisible(v[0], label_stride)
for v in scales]) // config.network.LABEL_STRIDE,
max([make_divisible(v[1], label_stride)
for v in scales]) // config.network.LABEL_STRIDE)),
('metric_label_' + str(scale_list[2]) + scale_name[2],
(config.TRAIN.BATCH_IMAGES, 9, 1,
max([make_divisible(v[0], label_stride)
for v in scales]) // config.network.LABEL_STRIDE,
max([make_divisible(v[1], label_stride)
for v in scales]) // config.network.LABEL_STRIDE))
]
else:
max_label_shape = [
('label',
(config.TRAIN.BATCH_IMAGES, 1,
max([make_divisible(v[0], label_stride)
for v in scales]) // config.network.LABEL_STRIDE,
max([make_divisible(v[1], label_stride)
for v in scales]) // config.network.LABEL_STRIDE)),
('metric_label',
(config.TRAIN.BATCH_IMAGES, 9, 1,
max([make_divisible(v[0], label_stride)
for v in scales]) // config.network.LABEL_STRIDE,
max([make_divisible(v[1], label_stride)
for v in scales]) // config.network.LABEL_STRIDE))
]
else:
if config.TRAIN.enable_crop:
max_label_shape = [('label',
(config.TRAIN.BATCH_IMAGES, 1,
config.TRAIN.crop_size[0] // label_stride,
config.TRAIN.crop_size[1] // label_stride))]
else:
max_label_shape = [
('label',
(config.TRAIN.BATCH_IMAGES, 1,
max([make_divisible(v[0], label_stride)
for v in scales]) // config.network.LABEL_STRIDE,
max([make_divisible(v[1], label_stride)
for v in scales]) // config.network.LABEL_STRIDE))
]
print "max_label_shapes", max_label_shape
if config.TRAIN.use_dynamic:
sym = sym_gen([max_data_shape])
# setup multi-gpu
input_batch_size = config.TRAIN.BATCH_IMAGES * len(ctx)
NUM_GPUS = len(ctx)
# 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=True) for image_set in image_sets
]
segdb = merge_segdb(segdbs)
# load training data
train_data = TrainDataLoader(sym,
segdb,
config,
batch_size=input_batch_size,
shuffle=config.TRAIN.SHUFFLE,
ctx=ctx,
use_context=config.network.use_context,
use_mult_label=config.network.use_mult_label,
use_metric=config.network.use_metric)
# loading val data
if config.TRAIN.eval_data_frequency > 0:
val_image_set = config.dataset.test_image_set
val_root_path = config.dataset.root_path
val_dataset = config.dataset.dataset
val_dataset_path = config.dataset.dataset_path
val_imdb = eval(val_dataset)(val_image_set,
val_root_path,
val_dataset_path,
result_path=final_output_path)
val_segdb = val_imdb.gt_segdb()
val_data = TrainDataLoader(
sym,
val_segdb,
config,
batch_size=input_batch_size,
shuffle=config.TRAIN.SHUFFLE,
ctx=ctx,
use_context=config.network.use_context,
use_mult_label=config.network.use_mult_label,
use_metric=config.network.use_metric)
else:
val_data = None
# print sym.list_arguments()
print 'providing maximum shape', max_data_shape, max_label_shape
max_data_shape, max_label_shape = train_data.infer_shape(
max_data_shape, max_label_shape)
print 'providing maximum shape', max_data_shape, max_label_shape
# infer shape
data_shape_dict = dict(train_data.provide_data_single +
train_data.provide_label_single)
if config.TRAIN.use_dynamic:
sym = sym_gen([train_data.provide_data_single])
pprint.pprint(data_shape_dict)
sym_instance.infer_shape(data_shape_dict)
nset = set()
for nm in sym.list_arguments():
if nm in nset:
raise ValueError('Duplicate names detected, %s' % str(nm))
nset.add(nm)
# 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)
sym_instance.check_parameter_shapes(arg_params,
aux_params,
data_shape_dict,
is_train=True)
preload_opt_states = load_preload_opt_states(prefix, begin_epoch)
# preload_opt_states = None
else:
print pretrained
arg_params, aux_params = load_param(pretrained, epoch, convert=True)
preload_opt_states = None
if not config.TRAIN.FINTUNE:
fixed_param_names = sym_instance.init_weights(
config, arg_params, aux_params)
sym_instance.check_parameter_shapes(arg_params,
aux_params,
data_shape_dict,
is_train=True)
# 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 xrange(NUM_GPUS)],
max_label_shapes=[max_label_shape for _ in xrange(NUM_GPUS)],
fixed_param_prefix=fixed_param_prefix)
# metric
imagecrossentropylossmetric = metric.ImageCrossEntropyLossMetric()
localmetric = metric.LocalImageCrossEntropyLossMetric()
globalmetric = metric.GlobalImageCrossEntropyLossMetric()
pixcelAccMetric = metric.PixcelAccMetric()
eval_metrics = mx.metric.CompositeEvalMetric()
if config.network.use_mult_label:
metric_list = [
imagecrossentropylossmetric, localmetric, globalmetric,
pixcelAccMetric
]
elif config.network.use_metric:
if config.TRAIN.use_crl_ses:
metric_list = [
imagecrossentropylossmetric,
metric.SigmoidPixcelAccMetric(1),
metric.SigmoidPixcelAccMetric(2),
metric.SigmoidPixcelAccMetric(3),
metric.CenterLossMetric(4),
metric.CenterLossMetric(5),
metric.CenterLossMetric(6), pixcelAccMetric
]
elif config.network.use_sigmoid_metric:
if config.TRAIN.use_mult_metric:
metric_list = [
imagecrossentropylossmetric,
metric.SigmoidPixcelAccMetric(1),
metric.SigmoidPixcelAccMetric(2),
metric.SigmoidPixcelAccMetric(3), pixcelAccMetric
]
else:
metric_list = [
imagecrossentropylossmetric,
metric.SigmoidPixcelAccMetric(), pixcelAccMetric
]
else:
if config.TRAIN.use_mult_metric:
metric_list = [
imagecrossentropylossmetric,
metric.MetricLossMetric(1),
metric.MetricLossMetric(2),
metric.MetricLossMetric(3), pixcelAccMetric
]
else:
metric_list = [
imagecrossentropylossmetric,
metric.MetricLossMetric(1), pixcelAccMetric
]
elif config.network.mult_loss:
metric_list = [
imagecrossentropylossmetric,
metric.MImageCrossEntropyLossMetric(1),
metric.MImageCrossEntropyLossMetric(2), pixcelAccMetric
]
elif config.TRAIN.use_center:
if config.TRAIN.use_one_center:
metric_list = [
imagecrossentropylossmetric, pixcelAccMetric,
metric.CenterLossMetric(1)
]
else:
metric_list = [
imagecrossentropylossmetric, pixcelAccMetric,
metric.CenterLossMetric(1),
metric.CenterLossMetric(2),
metric.CenterLossMetric(3)
]
else:
metric_list = [imagecrossentropylossmetric, pixcelAccMetric]
# rpn_eval_metric, rpn_cls_metric, rpn_bbox_metric, eval_metric, cls_metric, bbox_metric
for child_metric in metric_list:
eval_metrics.add(child_metric)
# callback
if False:
batch_end_callback = [
callback.Speedometer(train_data.batch_size,
frequent=args.frequent),
callback.TensorboardCallback(tensorboard_path,
prefix="train/batch")
]
epoch_end_callback = mx.callback.module_checkpoint(
mod, prefix, period=1, save_optimizer_states=True)
shared_tensorboard = batch_end_callback[1]
epoch_end_metric_callback = callback.TensorboardCallback(
tensorboard_path,
shared_tensorboard=shared_tensorboard,
prefix="train/epoch")
eval_end_callback = callback.TensorboardCallback(
tensorboard_path,
shared_tensorboard=shared_tensorboard,
prefix="val/epoch")
lr_callback = callback.LrCallback(
tensorboard_path,
shared_tensorboard=shared_tensorboard,
prefix='train/batch')
else:
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)
epoch_end_metric_callback = None
eval_end_callback = None
#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) / input_batch_size) for epoch in lr_epoch_diff
]
print 'lr', lr, 'lr_epoch_diff', lr_epoch_diff, 'lr_iters', lr_iters
if config.TRAIN.lr_type == "MultiStage":
lr_scheduler = LinearWarmupMultiStageScheduler(
lr_iters,
lr_factor,
config.TRAIN.warmup,
config.TRAIN.warmup_lr,
config.TRAIN.warmup_step,
args.frequent,
stop_lr=lr * 0.01)
elif config.TRAIN.lr_type == "MultiFactor":
lr_scheduler = LinearWarmupMultiFactorScheduler(
lr_iters, lr_factor, config.TRAIN.warmup, config.TRAIN.warmup_lr,
config.TRAIN.warmup_step, args.frequent)
if config.TRAIN.optimizer == "sgd":
optimizer_params = {
'momentum': config.TRAIN.momentum,
'wd': config.TRAIN.wd,
'learning_rate': lr,
'lr_scheduler': lr_scheduler,
'rescale_grad': 1.0,
'clip_gradient': None
}
optimizer = SGD(**optimizer_params)
elif config.TRAIN.optimizer == "adam":
optimizer_params = {
'learning_rate': lr,
'lr_scheduler': lr_scheduler,
'rescale_grad': 1.0,
'clip_gradient': None
}
optimizer = Adam(**optimizer_params)
print "optimizer adam"
freeze_layer_pattern = config.TRAIN.FIXED_PARAMS_PATTERN
if freeze_layer_pattern.strip():
args_lr_mult = {}
re_prog = re.compile(freeze_layer_pattern)
if freeze_layer_pattern:
fixed_param_names = [
name for name in sym.list_arguments() if re_prog.match(name)
]
print "============================"
print "fixed_params_names:"
print(fixed_param_names)
for name in fixed_param_names:
args_lr_mult[name] = config.TRAIN.FIXED_PARAMS_PATTERN_LR_MULT
print "============================"
else:
args_lr_mult = {}
optimizer.set_lr_mult(args_lr_mult)
# data_shape_dict = dict(train_data.provide_data_single + train_data.provide_label_single)
# if config.TRAIN.use_dynamic:
# sym = sym_gen([train_data.provide_data_single])
# pprint.pprint(data_shape_dict)
# sym_instance.infer_shape(data_shape_dict)
if not isinstance(train_data, PrefetchingIter) and config.TRAIN.use_thread:
train_data = PrefetchingIter(train_data)
if val_data:
if not isinstance(val_data, PrefetchingIter):
val_data = PrefetchingIter(val_data)
if Debug:
monitor = mx.monitor.Monitor(1)
else:
monitor = None
# 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,
eval_end_callback=eval_end_callback,
epoch_end_metric_callback=epoch_end_metric_callback,
optimizer=optimizer,
eval_data=val_data,
arg_params=arg_params,
aux_params=aux_params,
begin_epoch=begin_epoch,
num_epoch=end_epoch,
allow_missing=begin_epoch == 0,
allow_extra=True,
monitor=monitor,
preload_opt_states=preload_opt_states,
eval_data_frequency=config.TRAIN.eval_data_frequency)
def train_net(args, ctx, pretrained, epoch, prefix, begin_epoch, end_epoch, lr,
lr_step):
logger, final_output_path, _, tensorboard_path = create_env(
config.output_path, args.cfg, config.dataset.image_set)
prefix = os.path.join(final_output_path, prefix)
# print config
pprint.pprint(config)
logger.info('training config:{}\n'.format(pprint.pformat(config)))
print "config.symbol", config.symbol
sym_instance = eval(config.symbol + '.' + config.symbol)()
sym = sym_instance.get_symbol(config, is_train=True)
# setup multi-gpu
input_batch_size = config.TRAIN.BATCH_IMAGES * len(ctx)
NUM_GPUS = len(ctx)
# 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)
for image_set in image_sets
]
segdb = merge_segdb(segdbs)
# load training data
train_data = TrainDataLoader(sym,
segdb,
config,
batch_size=input_batch_size,
shuffle=config.TRAIN.SHUFFLE,
ctx=ctx)
# loading val data
val_image_set = config.dataset.test_image_set
val_root_path = config.dataset.root_path
val_dataset = config.dataset.dataset
val_dataset_path = config.dataset.dataset_path
val_imdb = eval(val_dataset)(val_image_set,
val_root_path,
val_dataset_path,
result_path=final_output_path)
val_segdb = val_imdb.gt_segdb()
val_data = TrainDataLoader(sym,
val_segdb,
config,
batch_size=input_batch_size,
shuffle=config.TRAIN.SHUFFLE,
ctx=ctx)
# infer max shape
max_scale = [(config.TRAIN.crop_size[0], config.TRAIN.crop_size[1])]
max_data_shape = [('data', (config.TRAIN.BATCH_IMAGES, 3,
max([v[0] for v in max_scale]),
max([v[1] for v in max_scale])))]
max_label_shape = [
('label',
(config.TRAIN.BATCH_IMAGES, 1,
max([v[0] for v in max_scale]) // config.network.LABEL_STRIDE,
max([v[1] for v in max_scale]) // config.network.LABEL_STRIDE))
]
max_data_shape, max_label_shape = train_data.infer_shape(
max_data_shape, max_label_shape)
print 'providing maximum shape', max_data_shape, max_label_shape
# infer 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)
preload_opt_states = load_preload_opt_states(prefix, begin_epoch)
else:
print pretrained
arg_params, aux_params = load_param(pretrained, epoch, convert=True)
preload_opt_states = None
sym_instance.init_weights(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 xrange(NUM_GPUS)],
max_label_shapes=[max_label_shape for _ in xrange(NUM_GPUS)],
fixed_param_prefix=fixed_param_prefix)
# metric
imagecrossentropylossmetric = metric.ImageCrossEntropyLossMetric()
pixcelAccMetric = metric.PixcelAccMetric()
eval_metrics = mx.metric.CompositeEvalMetric()
# rpn_eval_metric, rpn_cls_metric, rpn_bbox_metric, eval_metric, cls_metric, bbox_metric
for child_metric in [imagecrossentropylossmetric, pixcelAccMetric]:
eval_metrics.add(child_metric)
# callback
batch_end_callback = [
callback.Speedometer(train_data.batch_size, frequent=args.frequent),
callback.TensorboardCallback(tensorboard_path, prefix="train/batch")
]
epoch_end_callback = mx.callback.module_checkpoint(
mod, prefix, period=1, save_optimizer_states=True)
shared_tensorboard = batch_end_callback[1]
epoch_end_metric_callback = callback.TensorboardCallback(
tensorboard_path,
shared_tensorboard=shared_tensorboard,
prefix="train/epoch")
eval_end_callback = callback.TensorboardCallback(
tensorboard_path,
shared_tensorboard=shared_tensorboard,
prefix="val/epoch")
lr_callback = callback.LrCallback(tensorboard_path,
shared_tensorboard=shared_tensorboard,
prefix='train/batch')
#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) / input_batch_size) for epoch in lr_epoch_diff
]
print 'lr', lr, 'lr_epoch_diff', lr_epoch_diff, 'lr_iters', lr_iters
if config.TRAIN.lr_type == "MultiStage":
lr_scheduler = LinearWarmupMultiStageScheduler(
lr_iters,
lr_factor,
config.TRAIN.warmup,
config.TRAIN.warmup_lr,
config.TRAIN.warmup_step,
args.frequent,
stop_lr=lr * 0.01)
elif config.TRAIN.lr_type == "MultiFactor":
lr_scheduler = LinearWarmupMultiFactorScheduler(
lr_iters, lr_factor, config.TRAIN.warmup, config.TRAIN.warmup_lr,
config.TRAIN.warmup_step, args.frequent)
optimizer_params = {
'momentum': config.TRAIN.momentum,
'wd': config.TRAIN.wd,
'learning_rate': lr,
'lr_scheduler': lr_scheduler,
'rescale_grad': 1.0,
'clip_gradient': None
}
optimizer = SGD(**optimizer_params)
freeze_layer_pattern = config.TRAIN.FIXED_PARAMS_PATTERN
if freeze_layer_pattern.strip():
args_lr_mult = {}
re_prog = re.compile(freeze_layer_pattern)
fixed_param_names = [
name for name in sym.list_arguments() if re_prog.match(name)
]
print "fixed_params_names:"
print(fixed_param_names)
for name in fixed_param_names:
args_lr_mult[name] = config.TRAIN.FIXED_PARAMS_PATTERN_LR_MULT
else:
args_lr_mult = {}
optimizer.set_lr_mult(args_lr_mult)
if not isinstance(train_data, PrefetchingIter):
train_data = PrefetchingIter(train_data)
if not isinstance(val_data, PrefetchingIter):
val_data = PrefetchingIter(val_data)
if Debug:
monitor = mx.monitor.Monitor(1)
else:
monitor = None
initializer = mx.initializer.Xavier(magnitude=1, rnd_type="gaussian")
# 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,
eval_end_callback=eval_end_callback,
epoch_end_metric_callback=epoch_end_metric_callback,
optimizer=optimizer,
eval_data=val_data,
arg_params=arg_params,
aux_params=aux_params,
begin_epoch=begin_epoch,
num_epoch=end_epoch,
allow_missing=begin_epoch == 0,
allow_extra=True,
monitor=monitor,
preload_opt_states=preload_opt_states,
eval_data_frequency=config.TRAIN.eval_data_frequency,
initializer=initializer)