def main(unused_argv): assert FLAGS.checkpoint_dir, '`checkpoint_dir` is missing.' assert FLAGS.eval_dir, '`eval_dir` is missing.' if FLAGS.pipeline_config_path: configs = config_util.get_configs_from_pipeline_file( FLAGS.pipeline_config_path) else: configs = config_util.get_configs_from_multiple_files( model_config_path=FLAGS.model_config_path, eval_config_path=FLAGS.eval_config_path, eval_input_config_path=FLAGS.input_config_path) pipeline_proto = config_util.create_pipeline_proto_from_configs(configs) config_text = text_format.MessageToString(pipeline_proto) tf.gfile.MakeDirs(FLAGS.eval_dir) with tf.gfile.Open(os.path.join(FLAGS.eval_dir, 'pipeline.config'), 'wb') as f: f.write(config_text) model_config = configs['model'] lstm_config = configs['lstm_model'] eval_config = configs['eval_config'] input_config = configs['eval_input_config'] if FLAGS.eval_training_data: input_config.external_input_reader.CopyFrom( configs['train_input_config'].external_input_reader) lstm_config.eval_unroll_length = lstm_config.train_unroll_length model_fn = functools.partial( model_builder.build, model_config=model_config, lstm_config=lstm_config, is_training=False) def get_next(config, model_config, lstm_config, unroll_length): return seq_dataset_builder.build(config, model_config, lstm_config, unroll_length) create_input_dict_fn = functools.partial(get_next, input_config, model_config, lstm_config, lstm_config.eval_unroll_length) label_map = label_map_util.load_labelmap(input_config.label_map_path) max_num_classes = max([item.id for item in label_map.item]) categories = label_map_util.convert_label_map_to_categories( label_map, max_num_classes) if FLAGS.run_once: eval_config.max_evals = 1 evaluator.evaluate(create_input_dict_fn, model_fn, eval_config, categories, FLAGS.checkpoint_dir, FLAGS.eval_dir)
def main(unused_argv): assert FLAGS.checkpoint_dir, '`checkpoint_dir` is missing.' assert FLAGS.eval_dir, '`eval_dir` is missing.' if FLAGS.pipeline_config_path: configs = config_util.get_configs_from_pipeline_file( FLAGS.pipeline_config_path) else: configs = config_util.get_configs_from_multiple_files( model_config_path=FLAGS.model_config_path, eval_config_path=FLAGS.eval_config_path, eval_input_config_path=FLAGS.input_config_path) pipeline_proto = config_util.create_pipeline_proto_from_configs(configs) config_text = text_format.MessageToString(pipeline_proto) tf.gfile.MakeDirs(FLAGS.eval_dir) with tf.gfile.Open(os.path.join(FLAGS.eval_dir, 'pipeline.config'), 'wb') as f: f.write(config_text) model_config = configs['model'] lstm_config = configs['lstm_model'] eval_config = configs['eval_config'] input_config = configs['eval_input_config'] if FLAGS.eval_training_data: input_config.external_input_reader.CopyFrom( configs['train_input_config'].external_input_reader) lstm_config.eval_unroll_length = lstm_config.train_unroll_length model_fn = functools.partial( model_builder.build, model_config=model_config, lstm_config=lstm_config, is_training=False) def get_next(config, model_config, lstm_config, unroll_length): return seq_dataset_builder.build(config, model_config, lstm_config, unroll_length) create_input_dict_fn = functools.partial(get_next, input_config, model_config, lstm_config, lstm_config.eval_unroll_length) label_map = label_map_util.load_labelmap(input_config.label_map_path) max_num_classes = max([item.id for item in label_map.item]) categories = label_map_util.convert_label_map_to_categories( label_map, max_num_classes) if FLAGS.run_once: eval_config.max_evals = 1 evaluator.evaluate(create_input_dict_fn, model_fn, eval_config, categories, FLAGS.checkpoint_dir, FLAGS.eval_dir)
def main(_): assert FLAGS.train_dir, '`train_dir` is missing.' if FLAGS.task == 0: tf.gfile.MakeDirs(FLAGS.train_dir) if FLAGS.pipeline_config_path: configs = config_util.get_configs_from_pipeline_file( FLAGS.pipeline_config_path) if FLAGS.task == 0: tf.gfile.Copy(FLAGS.pipeline_config_path, os.path.join(FLAGS.train_dir, 'pipeline.config'), overwrite=True) else: configs = config_util.get_configs_from_multiple_files( model_config_path=FLAGS.model_config_path, train_config_path=FLAGS.train_config_path, train_input_config_path=FLAGS.input_config_path) if FLAGS.task == 0: for name, config in [('model.config', FLAGS.model_config_path), ('train.config', FLAGS.train_config_path), ('input.config', FLAGS.input_config_path)]: tf.gfile.Copy(config, os.path.join(FLAGS.train_dir, name), overwrite=True) model_config = configs['model'] lstm_config = configs['lstm_model'] train_config = configs['train_config'] input_config = configs['train_input_config'] model_fn = functools.partial(model_builder.build, model_config=model_config, lstm_config=lstm_config, is_training=True) def get_next(config, model_config, lstm_config, unroll_length): data_augmentation_options = [ preprocessor_builder.build(step) for step in train_config.data_augmentation_options ] return seq_dataset_builder.build(config, model_config, lstm_config, unroll_length, data_augmentation_options, batch_size=train_config.batch_size) create_input_dict_fn = functools.partial(get_next, input_config, model_config, lstm_config, lstm_config.train_unroll_length) env = json.loads(os.environ.get('TF_CONFIG', '{}')) cluster_data = env.get('cluster', None) cluster = tf.train.ClusterSpec(cluster_data) if cluster_data else None task_data = env.get('task', None) or {'type': 'master', 'index': 0} task_info = type('TaskSpec', (object, ), task_data) # Parameters for a single worker. ps_tasks = 0 worker_replicas = 1 worker_job_name = 'lonely_worker' task = 0 is_chief = True master = '' if cluster_data and 'worker' in cluster_data: # Number of total worker replicas include "worker"s and the "master". worker_replicas = len(cluster_data['worker']) + 1 if cluster_data and 'ps' in cluster_data: ps_tasks = len(cluster_data['ps']) if worker_replicas > 1 and ps_tasks < 1: raise ValueError( 'At least 1 ps task is needed for distributed training.') if worker_replicas >= 1 and ps_tasks > 0: # Set up distributed training. server = tf.train.Server(tf.train.ClusterSpec(cluster), protocol='grpc', job_name=task_info.type, task_index=task_info.index) if task_info.type == 'ps': server.join() return worker_job_name = '%s/task:%d' % (task_info.type, task_info.index) task = task_info.index is_chief = (task_info.type == 'master') master = server.target trainer.train(create_input_dict_fn, model_fn, train_config, master, task, FLAGS.num_clones, worker_replicas, FLAGS.clone_on_cpu, ps_tasks, worker_job_name, is_chief, FLAGS.train_dir)
def main(_): assert FLAGS.train_dir, '`train_dir` is missing.' if FLAGS.task == 0: tf.gfile.MakeDirs(FLAGS.train_dir) if FLAGS.pipeline_config_path: configs = config_util.get_configs_from_pipeline_file( FLAGS.pipeline_config_path) if FLAGS.task == 0: tf.gfile.Copy( FLAGS.pipeline_config_path, os.path.join(FLAGS.train_dir, 'pipeline.config'), overwrite=True) else: configs = config_util.get_configs_from_multiple_files( model_config_path=FLAGS.model_config_path, train_config_path=FLAGS.train_config_path, train_input_config_path=FLAGS.input_config_path) if FLAGS.task == 0: for name, config in [('model.config', FLAGS.model_config_path), ('train.config', FLAGS.train_config_path), ('input.config', FLAGS.input_config_path)]: tf.gfile.Copy( config, os.path.join(FLAGS.train_dir, name), overwrite=True) model_config = configs['model'] lstm_config = configs['lstm_model'] train_config = configs['train_config'] input_config = configs['train_input_config'] model_fn = functools.partial( model_builder.build, model_config=model_config, lstm_config=lstm_config, is_training=True) def get_next(config, model_config, lstm_config, unroll_length): data_augmentation_options = [ preprocessor_builder.build(step) for step in train_config.data_augmentation_options ] return seq_dataset_builder.build( config, model_config, lstm_config, unroll_length, data_augmentation_options, batch_size=train_config.batch_size) create_input_dict_fn = functools.partial(get_next, input_config, model_config, lstm_config, lstm_config.train_unroll_length) env = json.loads(os.environ.get('TF_CONFIG', '{}')) cluster_data = env.get('cluster', None) cluster = tf.train.ClusterSpec(cluster_data) if cluster_data else None task_data = env.get('task', None) or {'type': 'master', 'index': 0} task_info = type('TaskSpec', (object,), task_data) # Parameters for a single worker. ps_tasks = 0 worker_replicas = 1 worker_job_name = 'lonely_worker' task = 0 is_chief = True master = '' if cluster_data and 'worker' in cluster_data: # Number of total worker replicas include "worker"s and the "master". worker_replicas = len(cluster_data['worker']) + 1 if cluster_data and 'ps' in cluster_data: ps_tasks = len(cluster_data['ps']) if worker_replicas > 1 and ps_tasks < 1: raise ValueError('At least 1 ps task is needed for distributed training.') if worker_replicas >= 1 and ps_tasks > 0: # Set up distributed training. server = tf.train.Server( tf.train.ClusterSpec(cluster), protocol='grpc', job_name=task_info.type, task_index=task_info.index) if task_info.type == 'ps': server.join() return worker_job_name = '%s/task:%d' % (task_info.type, task_info.index) task = task_info.index is_chief = (task_info.type == 'master') master = server.target trainer.train(create_input_dict_fn, model_fn, train_config, master, task, FLAGS.num_clones, worker_replicas, FLAGS.clone_on_cpu, ps_tasks, worker_job_name, is_chief, FLAGS.train_dir)