import os import copy import tensorflow as tf from tensorflow.python.client import device_lib from e2eflow.core.train import Trainer from e2eflow.experiment import Experiment from e2eflow.util import convert_input_strings from e2eflow.kitti.input import KITTIInput from e2eflow.kitti.data import KITTIData from e2eflow.chairs.data import ChairsData from e2eflow.chairs.input import ChairsInput from e2eflow.sintel.data import SintelData from e2eflow.sintel.input import SintelInput from e2eflow.synthia.data import SynthiaData from e2eflow.cityscapes.data import CityscapesData tf.app.flags.DEFINE_string('ex', 'default', 'Name of the experiment.' 'If the experiment folder already exists in the log dir, ' 'training will be continued from the latest checkpoint.') tf.app.flags.DEFINE_boolean('debug', False, 'Enable image summaries and disable checkpoint writing for debugging.') tf.app.flags.DEFINE_boolean('ow', False, 'Overwrites a previous experiment with the same name (if present)' 'instead of attempting to continue from its latest checkpoint.') FLAGS = tf.app.flags.FLAGS def main(argv=None): experiment = Experiment( name=FLAGS.ex, overwrite=FLAGS.ow) dirs = experiment.config['dirs'] run_config = experiment.config['run'] gpu_list_param = run_config['gpu_list'] if isinstance(gpu_list_param, int): gpu_list = [gpu_list_param] os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_list_param) else: gpu_list = list(range(len(gpu_list_param.split(',')))) os.environ['CUDA_VISIBLE_DEVICES'] = gpu_list_param gpu_batch_size = int(run_config['batch_size'] / max(len(gpu_list), 1)) devices = ['/gpu:' + str(gpu_num) for gpu_num in gpu_list] train_dataset = run_config.get('dataset', 'kitti') kdata = KITTIData(data_dir=dirs['data'], fast_dir=dirs.get('fast'), stat_log_dir=None, development=run_config['development']) einput = KITTIInput(data=kdata, batch_size=1, normalize=False, dims=(384, 1280)) if train_dataset == 'chairs': cconfig = copy.deepcopy(experiment.config['train']) cconfig.update(experiment.config['train_chairs']) convert_input_strings(cconfig, dirs) citers = cconfig.get('num_iters', 0) cdata = ChairsData(data_dir=dirs['data'], fast_dir=dirs.get('fast'), stat_log_dir=None, development=run_config['development']) cinput = ChairsInput(data=cdata, batch_size=gpu_batch_size, normalize=False, dims=(cconfig['height'], cconfig['width'])) tr = Trainer( lambda shift: cinput.input_raw(swap_images=False, shift=shift * run_config['batch_size']), lambda: einput.input_train_2012(), params=cconfig, normalization=cinput.get_normalization(), train_summaries_dir=experiment.train_dir, eval_summaries_dir=experiment.eval_dir, experiment=FLAGS.ex, ckpt_dir=experiment.save_dir, debug=FLAGS.debug, interactive_plot=run_config.get('interactive_plot'), devices=devices) tr.run(0, citers) elif train_dataset == 'kitti': kconfig = copy.deepcopy(experiment.config['train']) kconfig.update(experiment.config['train_kitti']) convert_input_strings(kconfig, dirs) kiters = kconfig.get('num_iters', 0) kinput = KITTIInput(data=kdata, batch_size=gpu_batch_size, normalize=False, skipped_frames=True, dims=(kconfig['height'], kconfig['width'])) tr = Trainer( lambda shift: kinput.input_raw(swap_images=False, center_crop=True, shift=shift * run_config['batch_size']), lambda: einput.input_train_2012(), params=kconfig, normalization=kinput.get_normalization(), train_summaries_dir=experiment.train_dir, eval_summaries_dir=experiment.eval_dir, experiment=FLAGS.ex, ckpt_dir=experiment.save_dir, debug=FLAGS.debug, interactive_plot=run_config.get('interactive_plot'), devices=devices) tr.run(0, kiters) elif train_dataset == 'cityscapes': kconfig = copy.deepcopy(experiment.config['train']) kconfig.update(experiment.config['train_cityscapes']) convert_input_strings(kconfig, dirs) kiters = kconfig.get('num_iters', 0) cdata = CityscapesData(data_dir=dirs['data'], fast_dir=dirs.get('fast'), stat_log_dir=None, development=run_config['development']) kinput = KITTIInput(data=cdata, batch_size=gpu_batch_size, normalize=False, skipped_frames=False, dims=(kconfig['height'], kconfig['width'])) tr = Trainer( lambda shift: kinput.input_raw(swap_images=False, center_crop=True, skip=[0, 1], shift=shift * run_config['batch_size']), lambda: einput.input_train_2012(), params=kconfig, normalization=kinput.get_normalization(), train_summaries_dir=experiment.train_dir, eval_summaries_dir=experiment.eval_dir, experiment=FLAGS.ex, ckpt_dir=experiment.save_dir, debug=FLAGS.debug, interactive_plot=run_config.get('interactive_plot'), devices=devices) tr.run(0, kiters) elif train_dataset == 'synthia': sconfig = copy.deepcopy(experiment.config['train']) sconfig.update(experiment.config['train_synthia']) convert_input_strings(sconfig, dirs) siters = sconfig.get('num_iters', 0) sdata = SynthiaData(data_dir=dirs['data'], fast_dir=dirs.get('fast'), stat_log_dir=None, development=run_config['development']) sinput = KITTIInput(data=sdata, batch_size=gpu_batch_size, normalize=False, dims=(sconfig['height'], sconfig['width'])) tr = Trainer( lambda shift: sinput.input_raw(swap_images=False, shift=shift * run_config['batch_size']), lambda: einput.input_train_2012(), params=sconfig, normalization=sinput.get_normalization(), train_summaries_dir=experiment.train_dir, eval_summaries_dir=experiment.eval_dir, experiment=FLAGS.ex, ckpt_dir=experiment.save_dir, debug=FLAGS.debug, interactive_plot=run_config.get('interactive_plot'), devices=devices) tr.run(0, siters) elif train_dataset == 'kitti_ft': ftconfig = copy.deepcopy(experiment.config['train']) ftconfig.update(experiment.config['train_kitti_ft']) convert_input_strings(ftconfig, dirs) ftiters = ftconfig.get('num_iters', 0) ftinput = KITTIInput(data=kdata, batch_size=gpu_batch_size, normalize=False, dims=(ftconfig['height'], ftconfig['width'])) tr = Trainer( lambda shift: ftinput.input_train_gt(40), lambda: einput.input_train_2015(40), supervised=True, params=ftconfig, normalization=ftinput.get_normalization(), train_summaries_dir=experiment.train_dir, eval_summaries_dir=experiment.eval_dir, experiment=FLAGS.ex, ckpt_dir=experiment.save_dir, debug=FLAGS.debug, interactive_plot=run_config.get('interactive_plot'), devices=devices) tr.run(0, ftiters) else: raise ValueError( "Invalid dataset. Dataset must be one of " "{synthia, kitti, kitti_ft, cityscapes, chairs}") if not FLAGS.debug: experiment.conclude() if __name__ == '__main__': tf.app.run()