def train_net(args, ctx):
logger.auto_set_dir()
from symbols.gluon_deeplabv2 import resnet101_deeplab_new
input_batch_size = args.batch_size * len(ctx)
sym_instance = resnet101_deeplab_new()
sym = sym_instance.get_symbol(NUM_CLASSES, is_train=True)
eval_sym_instance = resnet101_deeplab_new()
eval_sym = eval_sym_instance.get_symbol(NUM_CLASSES, is_train=False)
train_data = get_data("train_aug", DATA_DIR, LIST_DIR, len(ctx))
eval_data = get_data("val", DATA_DIR, LIST_DIR, len(ctx))
# infer shape
data_shape_dict = {'data':(args.batch_size, 3, args.crop_size[0],args.crop_size[1])
,'label':(args.batch_size, 1, args.crop_size[0],args.crop_size[1])}
sym_instance.infer_shape(data_shape_dict)
# load and initialize params
begin_epoch = 1
mod = MutableModule(sym, data_names=['data'], label_names=['label'],context=ctx, fixed_param_prefix=fixed_param_prefix)
# metric
fcn_loss_metric = metric.FCNLogLossMetric(args.frequent,PascalVOC12.class_num())
eval_metrics = mx.metric.CompositeEvalMetric()
eval_metrics.add(fcn_loss_metric)
batch_end_callbacks = [callback.Speedometer(input_batch_size, frequent=args.frequent)]
epoch_end_callbacks = \
[mx.callback.module_checkpoint(mod, os.path.join(logger.get_logger_dir(),"mxnetgo"), period=1, save_optimizer_states=True),
]
lr_scheduler = StepScheduler(train_data.size()*EPOCH_SCALE,lr_step_list)
# optimizer
optimizer_params = {'momentum': 0.9,
'wd': 0.0005,
'learning_rate': 2.5e-4,
'lr_scheduler': lr_scheduler,
'rescale_grad': 1.0,
'clip_gradient': None}
logger.info("epoch scale = {}".format(EPOCH_SCALE))
mod.fit(train_data=train_data, args = args, eval_sym=eval_sym, eval_sym_instance=eval_sym_instance, eval_data=eval_data, eval_metric=eval_metrics, epoch_end_callback=epoch_end_callbacks,
batch_end_callback=batch_end_callbacks, kvstore=kvstore,
optimizer='sgd', optimizer_params=optimizer_params,
arg_params=None, aux_params=None, begin_epoch=begin_epoch, num_epoch=end_epoch,epoch_scale=EPOCH_SCALE, validation_on_last=validation_on_last)
os.environ['PYTHONUNBUFFERED'] = '1'
os.environ['MXNET_CUDNN_AUTOTUNE_DEFAULT'] = '0'
os.environ['MXNET_ENABLE_GPU_P2P'] = '0'
IGNORE_LABEL = 255
CROP_HEIGHT = 473
CROP_WIDTH = 473
tile_height = 321
tile_width = 321
batch_size = 10
EPOCH_SCALE = 4
end_epoch = 9
lr_step_list = [(6, 1e-3), (9, 1e-4)]
NUM_CLASSES = PascalVOC12.class_num()
validation_on_last = end_epoch
kvstore = "device"
fixed_param_prefix = ['conv0_weight', 'stage1', 'stage2', 'stage3']
symbol_str = "symbol_resnet_deeplabv1"
def parse_args():
parser = argparse.ArgumentParser(description='Train deeplab network')
# training
parser.add_argument("--gpu", default="5")
parser.add_argument('--frequent',
help='frequency of logging',
default=200,
type=int)
def train_net(args, ctx):
logger.auto_set_dir()
sym_instance = resnet101_deeplab_new()
sym = sym_instance.get_symbol(NUM_CLASSES,
is_train=True,
use_global_stats=False)
eval_sym_instance = resnet101_deeplab_new()
eval_sym = eval_sym_instance.get_symbol(NUM_CLASSES,
is_train=False,
use_global_stats=True)
# setup multi-gpu
gpu_nums = len(ctx)
input_batch_size = args.batch_size * gpu_nums
train_data = get_data("train_aug", DATA_DIR, LIST_DIR, len(ctx))
test_data = get_data("val", DATA_DIR, LIST_DIR, len(ctx))
# infer shape
data_shape_dict = {
'data': (args.batch_size, 3, args.crop_size[0], args.crop_size[1]),
'label': (args.batch_size, 1, args.crop_size[0], args.crop_size[1])
}
pprint.pprint(data_shape_dict)
sym_instance.infer_shape(data_shape_dict)
# load and initialize params
epoch_string = args.load.rsplit("-", 2)[1]
begin_epoch = 1
if not args.scratch:
begin_epoch = int(epoch_string)
logger.info('continue training from {}'.format(begin_epoch))
arg_params, aux_params = load_init_param(args.load, convert=True)
else:
logger.info(args.load)
arg_params, aux_params = load_init_param(args.load, convert=True)
sym_instance.init_weights(arg_params, aux_params)
# check parameter shapes
sym_instance.check_parameter_shapes(arg_params, aux_params,
data_shape_dict)
data_names = ['data']
label_names = ['label']
mod = MutableModule(sym,
data_names=data_names,
label_names=label_names,
context=ctx,
fixed_param_prefix=fixed_param_prefix)
# decide training params
# metric
fcn_loss_metric = metric.FCNLogLossMetric(args.frequent,
PascalVOC12.class_num())
eval_metrics = mx.metric.CompositeEvalMetric()
eval_metrics.add(fcn_loss_metric)
# callback
batch_end_callbacks = [
callback.Speedometer(input_batch_size, frequent=args.frequent)
]
#batch_end_callbacks = [mx.callback.ProgressBar(total=train_data.size/train_data.batch_size)]
epoch_end_callbacks = \
[mx.callback.module_checkpoint(mod, os.path.join(logger.get_logger_dir(),"mxnetgo"), period=1, save_optimizer_states=True),
]
lr_scheduler = StepScheduler(train_data.size() * EPOCH_SCALE, lr_step_list)
# optimizer
optimizer_params = {
'momentum': 0.9,
'wd': 0.0005,
'learning_rate': 2.5e-4,
'lr_scheduler': lr_scheduler,
'rescale_grad': 1.0,
'clip_gradient': None
}
logger.info("epoch scale = {}".format(EPOCH_SCALE))
mod.fit(train_data=train_data,
args=args,
eval_sym=eval_sym,
eval_sym_instance=eval_sym_instance,
eval_data=test_data,
eval_metric=eval_metrics,
epoch_end_callback=epoch_end_callbacks,
batch_end_callback=batch_end_callbacks,
kvstore=kvstore,
optimizer='sgd',
optimizer_params=optimizer_params,
arg_params=arg_params,
aux_params=aux_params,
begin_epoch=begin_epoch,
num_epoch=end_epoch,
epoch_scale=EPOCH_SCALE,
validation_on_last=validation_on_last)
