def _export(config, restore_path, image_path): if restore_path is None: restore_file = executor.search_restore_filename( environment.CHECKPOINTS_DIR) restore_path = os.path.join(environment.CHECKPOINTS_DIR, restore_file) print("Restore from {}".format(restore_path)) if not os.path.exists("{}.index".format(restore_path)): raise Exception("restore file {} dont exists.".format(restore_path)) output_root_dir = os.path.join(environment.EXPERIMENT_DIR, "export") output_root_dir = os.path.join(output_root_dir, os.path.basename(restore_path)) if not os.path.exists(output_root_dir): os.makedirs(output_root_dir) graph = tf.Graph() ModelClass = config.NETWORK_CLASS network_kwargs = dict( (key.lower(), val) for key, val in config.NETWORK.items()) with graph.as_default(): model = ModelClass( classes=config.CLASSES, is_debug=config.IS_DEBUG, **network_kwargs, ) is_training = tf.constant(False, name="is_training") images_placeholder, _ = model.placeholders() model.inference(images_placeholder, is_training) init_op = tf.compat.v1.global_variables_initializer() saver = tf.compat.v1.train.Saver(max_to_keep=50) session_config = tf.compat.v1.ConfigProto() sess = tf.compat.v1.Session(graph=graph, config=session_config) sess.run(init_op) saver.restore(sess, restore_path) main_output_dir = os.path.join( output_root_dir, "{}x{}".format(config.IMAGE_SIZE[0], config.IMAGE_SIZE[1])) if not os.path.exists(main_output_dir): os.makedirs(main_output_dir) # save inference values as npy files for runtime inference test and debug. if image_path: all_ops = _minimal_operations(sess) inference_values_output_dir = os.path.join(main_output_dir, "inference_test_data") if not os.path.exists(inference_values_output_dir): os.makedirs(inference_values_output_dir) raw_image = load_image(image_path) image = _pre_process(raw_image, config.PRE_PROCESSOR, config.DATA_FORMAT) images = np.expand_dims(image, axis=0) feed_dict = { images_placeholder: images, } all_outputs = [] index = 0 for op in all_ops: for op_output in op.outputs: # HACK: This is for TensorFlow bug workaround. # We can remove following 4 lines once it's been resolved in TensorFlow # Issue link: https://github.com/tensorflow/tensorflow/issues/36456 if (not tf.config.experimental.list_physical_devices('GPU') and "FusedBatchNormV3" in op_output.name and int( op_output.name.split(":")[1]) in set(range(1, 6))): continue val = sess.run(op_output.name, feed_dict=feed_dict) name = '%03d' % index + '_' + op_output.name.replace('/', '_') all_outputs.append({'val': val, 'name': name}) index += 1 _save_all_operation_outputs(image_path, inference_values_output_dir, image, raw_image, all_outputs, config.IMAGE_SIZE) yaml_names = config_util.save_yaml(main_output_dir, config) pb_name = executor.save_pb_file(sess, main_output_dir) message = """ Create pb and yaml files in: {} pb: {} yaml: {}, {} """.format(main_output_dir, pb_name, *yaml_names) if image_path: message += "Create npy files in under `inference_test_data` folder \n" message += "npy: {}".format([d["name"] for d in all_outputs] + [ "raw_image", "preprocessed_image", ]) print(message) print("finish") return main_output_dir
def start_training(config): use_horovod = horovod_util.is_enabled() print("use_horovod:", use_horovod) if use_horovod: hvd = horovod_util.setup() rank = hvd.rank() local_rank = hvd.local_rank() else: rank = 0 local_rank = -1 ModelClass = config.NETWORK_CLASS network_kwargs = {key.lower(): val for key, val in config.NETWORK.items()} train_dataset = setup_dataset(config, "train", rank, local_rank) print("train dataset num:", train_dataset.num_per_epoch) validation_dataset = setup_dataset(config, "validation", rank, local_rank) print("validation dataset num:", validation_dataset.num_per_epoch) graph = tf.Graph() with graph.as_default(): if config.TASK == Tasks.OBJECT_DETECTION: model = ModelClass( classes=train_dataset.classes, num_max_boxes=train_dataset.num_max_boxes, is_debug=config.IS_DEBUG, **network_kwargs, ) else: model = ModelClass( classes=train_dataset.classes, is_debug=config.IS_DEBUG, **network_kwargs, ) is_training_placeholder = tf.compat.v1.placeholder(tf.bool, name="is_training_placeholder") images_placeholder, labels_placeholder = model.placeholders() output = model.inference(images_placeholder, is_training_placeholder) loss = model.loss(output, labels_placeholder) opt = model.optimizer() if use_horovod: # add Horovod Distributed Optimizer opt = hvd.DistributedOptimizer(opt) train_op = model.train(loss, opt) metrics_ops_dict, metrics_update_op = model.metrics(output, labels_placeholder) # TODO(wakisaka): Deal with many networks. model.summary(output, labels_placeholder) summary_op = tf.compat.v1.summary.merge_all() metrics_summary_op = executor.metrics_summary_op(metrics_ops_dict) init_op = tf.compat.v1.global_variables_initializer() reset_metrics_op = tf.compat.v1.local_variables_initializer() if use_horovod: # add Horovod broadcasting variables from rank 0 to all bcast_global_variables_op = hvd.broadcast_global_variables(0) saver = tf.compat.v1.train.Saver(max_to_keep=config.KEEP_CHECKPOINT_MAX) if config.IS_PRETRAIN: all_vars = tf.compat.v1.global_variables() pretrain_var_list = [ var for var in all_vars if var.name.startswith(tuple(config.PRETRAIN_VARS)) ] print("pretrain_vars", [ var.name for var in pretrain_var_list ]) pretrain_saver = tf.compat.v1.train.Saver(pretrain_var_list, name="pretrain_saver") if use_horovod: # For distributed training session_config = tf.compat.v1.ConfigProto( gpu_options=tf.compat.v1.GPUOptions( allow_growth=True, visible_device_list=str(hvd.local_rank()) ) ) else: # TODO(wakisaka): For debug. # session_config = tf.ConfigProto( # gpu_options=tf.GPUOptions( # allow_growth=True, # per_process_gpu_memory_fraction=0.1 # ) # ) session_config = tf.compat.v1.ConfigProto() # tf.ConfigProto(log_device_placement=True) # TODO(wakisaka): XLA JIT # session_config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1 sess = tf.compat.v1.Session(graph=graph, config=session_config) sess.run([init_op, reset_metrics_op]) executor.save_pb_file(sess, environment.CHECKPOINTS_DIR) if rank == 0: train_writer = tf.compat.v1.summary.FileWriter(environment.TENSORBOARD_DIR + "/train", sess.graph) val_writer = tf.compat.v1.summary.FileWriter(environment.TENSORBOARD_DIR + "/validation") if config.IS_PRETRAIN: print("------- Load pretrain data ----------") pretrain_saver.restore(sess, os.path.join(config.PRETRAIN_DIR, config.PRETRAIN_FILE)) # for recovery ckpt = tf.train.get_checkpoint_state(environment.CHECKPOINTS_DIR) if ckpt and ckpt.model_checkpoint_path: print("--------- Restore last checkpoint -------------") saver.restore(sess, ckpt.model_checkpoint_path) # saver.recover_last_checkpoints(ckpt.model_checkpoint_path) last_step = sess.run(model.global_step) # TODO(wakisaka): tensorflow v1.3 remain previous event log in tensorboard. # https://github.com/tensorflow/tensorflow/blob/r1.3/tensorflow/python/training/supervisor.py#L1072 train_writer.add_session_log(SessionLog(status=SessionLog.START), global_step=last_step + 1) val_writer.add_session_log(SessionLog(status=SessionLog.START), global_step=last_step + 1) print("recovered. last step", last_step) if use_horovod: # broadcast variables from rank 0 to all other processes sess.run(bcast_global_variables_op) last_step = sess.run(model.global_step) # Calculate max steps. The priority of config.MAX_EPOCHS is higher than config.MAX_STEPS. if "MAX_EPOCHS" in config: max_steps = int(train_dataset.num_per_epoch / config.BATCH_SIZE * config.MAX_EPOCHS) else: max_steps = config.MAX_STEPS progbar = Progbar(max_steps) if rank == 0: progbar.update(last_step) for step in range(last_step, max_steps): images, labels = train_dataset.feed() feed_dict = { is_training_placeholder: True, images_placeholder: images, labels_placeholder: labels, } if step * ((step + 1) % config.SUMMARISE_STEPS) == 0 and rank == 0: # Runtime statistics for develop. # run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE) # run_metadata = tf.RunMetadata() sess.run(reset_metrics_op) _, summary, _ = sess.run( [train_op, summary_op, metrics_update_op], feed_dict=feed_dict, # options=run_options, # run_metadata=run_metadata, ) # train_writer.add_run_metadata(run_metadata, "step: {}".format(step + 1)) train_writer.add_summary(summary, step + 1) metrics_summary = sess.run(metrics_summary_op) train_writer.add_summary(metrics_summary, step + 1) train_writer.flush() else: sess.run([train_op], feed_dict=feed_dict) to_be_saved = step == 0 or (step + 1) == max_steps or (step + 1) % config.SAVE_CHECKPOINT_STEPS == 0 if to_be_saved and rank == 0: _save_checkpoint(saver, sess, model.global_step) if step == 0 or (step + 1) % config.TEST_STEPS == 0: # init metrics values sess.run(reset_metrics_op) test_step_size = int(math.ceil(validation_dataset.num_per_epoch / config.BATCH_SIZE)) for test_step in range(test_step_size): images, labels = validation_dataset.feed() feed_dict = { is_training_placeholder: False, images_placeholder: images, labels_placeholder: labels, } if test_step % config.SUMMARISE_STEPS == 0: summary, _ = sess.run([summary_op, metrics_update_op], feed_dict=feed_dict) if rank == 0: val_writer.add_summary(summary, step + 1) val_writer.flush() else: sess.run([metrics_update_op], feed_dict=feed_dict) metrics_summary = sess.run(metrics_summary_op) if rank == 0: val_writer.add_summary(metrics_summary, step + 1) val_writer.flush() if rank == 0: progbar.update(step + 1) # training loop end. train_dataset.close() validation_dataset.close() print("Done")