def train_net(args, ctx): logger.auto_set_dir() sym_instance = resnet101_deeplab_new() sym = sym_instance.get_symbol(NUM_CLASSES, is_train=True,memonger=False) #digraph = mx.viz.plot_network(sym, save_format='pdf') #digraph.render() # setup multi-gpu gpu_nums = len(ctx) input_batch_size = args.batch_size * gpu_nums train_data = get_data("train", DATA_DIR, LIST_DIR, len(ctx)) test_data = get_data("val", DATA_DIR, LIST_DIR, len(ctx)) # infer max shape max_scale = [args.crop_size] max_data_shape = [('data', (args.batch_size, 3, max([v[0] for v in max_scale]), max([v[1] for v in max_scale])))] max_label_shape = [('label', (args.batch_size, 1, max([v[0] for v in max_scale]), max([v[1] for v in max_scale])))] # 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) eval_sym_instance = resnet101_deeplab_new() # 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, max_data_shapes=[max_data_shape for _ in xrange(gpu_nums)], max_label_shapes=[max_label_shape for _ in xrange(gpu_nums)], fixed_param_prefix=fixed_param_prefix) # decide training params # metric fcn_loss_metric = metric.FCNLogLossMetric(args.frequent) eval_metrics = mx.metric.CompositeEvalMetric() for child_metric in [fcn_loss_metric]: eval_metrics.add(child_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 = { 'wd': 0.0005, 'learning_rate': init_lr, '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_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='adam', 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)
def do_validation(ctx): #logger.auto_set_dir() test_data = get_data("val", DATA_DIR, LIST_DIR, len(ctx)) ctx = [mx.gpu(int(i)) for i in args.gpu.split(',')] sym_instance = resnet101_deeplab_new() val_provide_data = [[("data", (1, 3, tile_height, tile_width))]] val_provide_label = [[("softmax_label", (1, 1, tile_height, tile_width))]] data_shape_dict = { 'data': (1, 3, tile_height, tile_width), 'softmax_label': (1, 1, tile_height, tile_width) } eval_sym = sym_instance.get_symbol(NUM_CLASSES, is_train=False) sym_instance.infer_shape(data_shape_dict) arg_params, aux_params = load_init_param(args.load, process=True) sym_instance.check_parameter_shapes(arg_params, aux_params, data_shape_dict, is_train=False) data_names = ['data'] label_names = ['softmax_label'] # create predictor predictor = Predictor(eval_sym, data_names, label_names, context=ctx, provide_data=val_provide_data, provide_label=val_provide_label, arg_params=arg_params, aux_params=aux_params) logger.info("begin prediction.. evaluation size: {}x{}".format( tile_width, tile_height)) if args.vis: from mxnetgo.myutils.fs import mkdir_p vis_dir = os.path.join(logger.get_logger_dir(), "vis") mkdir_p(vis_dir) stats = MIoUStatistics(NUM_CLASSES) test_data.reset_state() nbatch = 0 for data, label in tqdm(test_data.get_data()): output_all = predict_scaler(data, predictor, scales=[0.9, 1.0, 1.1], classes=NUM_CLASSES, tile_size=(tile_height, tile_width), is_densecrf=False, nbatch=nbatch, val_provide_data=val_provide_data, val_provide_label=val_provide_label) output_all = np.argmax(output_all, axis=0) label = np.squeeze(label) if args.vis: cv2.imwrite(os.path.join(vis_dir, "{}.jpg".format(nbatch)), visualize_label(output_all)) stats.feed(output_all, label) # very time-consuming nbatch += 1 logger.info("mIoU: {}, meanAcc: {}, acc: {} ".format( stats.mIoU, stats.mean_accuracy, stats.accuracy))