def main(argv): del argv # Unused. params = factory.config_generator(FLAGS.model) if FLAGS.config_file: params = params_dict.override_params_dict(params, FLAGS.config_file, is_strict=True) params = params_dict.override_params_dict(params, FLAGS.params_override, is_strict=True) params.validate() params.lock() model_params = dict(params.as_dict(), use_tpu=FLAGS.use_tpu, mode=tf.estimator.ModeKeys.PREDICT, transpose_input=False) print(' - Setting up TPUEstimator...') estimator = tf.contrib.tpu.TPUEstimator( model_fn=serving.serving_model_fn_builder( FLAGS.use_tpu, FLAGS.output_image_info, FLAGS.output_normalized_coordinates, FLAGS.cast_num_detections_to_float), model_dir=None, config=tpu_config.RunConfig( tpu_config=tpu_config.TPUConfig(iterations_per_loop=1), master='local', evaluation_master='local'), params=model_params, use_tpu=FLAGS.use_tpu, train_batch_size=FLAGS.batch_size, predict_batch_size=FLAGS.batch_size, export_to_tpu=FLAGS.use_tpu, export_to_cpu=True) print(' - Exporting the model...') input_type = FLAGS.input_type image_size = [int(x) for x in FLAGS.input_image_size.split(',')] export_path = estimator.export_saved_model( export_dir_base=FLAGS.export_dir, serving_input_receiver_fn=functools.partial( serving.serving_input_fn, batch_size=FLAGS.batch_size, desired_image_size=image_size, stride=(2**params.anchor.max_level), input_type=input_type, input_name=FLAGS.input_name), checkpoint_path=FLAGS.checkpoint_path) print(' - Done! path: %s' % export_path)
def run(callbacks=None): # keras_utils.set_session_config(enable_xla=FLAGS.enable_xla) params = config_factory.config_generator(FLAGS.model) params = params_dict.override_params_dict( params, FLAGS.config_file, is_strict=True) params = params_dict.override_params_dict( params, FLAGS.params_override, is_strict=True) params.override( { 'strategy_type': FLAGS.strategy_type, 'model_dir': FLAGS.model_dir, }, is_strict=False) params.validate() params.lock() pp = pprint.PrettyPrinter() params_str = pp.pformat(params.as_dict()) logging.info('Model Parameters: {}'.format(params_str)) train_input_fn = None eval_input_fn = None training_file_pattern = FLAGS.training_file_pattern or params.train.train_file_pattern eval_file_pattern = FLAGS.eval_file_pattern or params.eval.eval_file_pattern if not training_file_pattern and not eval_file_pattern: raise ValueError('Must provide at least one of training_file_pattern and ' 'eval_file_pattern.') if training_file_pattern: # Use global batch size for single host. train_input_fn = input_reader.InputFn( file_pattern=training_file_pattern, params=params, mode=input_reader.ModeKeys.TRAIN, batch_size=params.train.batch_size) if eval_file_pattern: eval_input_fn = input_reader.InputFn( file_pattern=eval_file_pattern, params=params, mode=input_reader.ModeKeys.PREDICT_WITH_GT, batch_size=params.eval.batch_size, num_examples=params.eval.eval_samples) # estimator_run(params, train_input_fn) return run_executor( params, mode=ModeKeys.TRAIN, train_input_fn=train_input_fn, callbacks=callbacks)
def __init__( self, config_file: str, checkpoint_path: str, batch_size: int, resize_shape: tuple[int, int], cache_dir: str, device: int | None = None, ): self.device = device self.batch_size = batch_size self.resize_shape = resize_shape params = config_factory.config_generator("mask_rcnn") if config_file: params = params_dict.override_params_dict(params, config_file, is_strict=True) params.validate() params.lock() self.max_level = params.architecture.max_level self._model = model_factory.model_generator(params) estimator = tf.estimator.Estimator(model_fn=self._model_fn, ) # Use SavedModel instead of Estimator.predcit() # because it is difficult to download images from GCS # when executing these codes on Vertex Pipelines. with tempfile.TemporaryDirectory() as tmpdir: export_dir_parent = cache_dir or tmpdir children = list(Path(export_dir_parent).glob("*")) if children == []: logger.info(f"export saved_model: {export_dir_parent}") estimator.export_saved_model( export_dir_base=export_dir_parent, serving_input_receiver_fn=self._serving_input_receiver_fn, checkpoint_path=checkpoint_path, ) children = list(Path(export_dir_parent).glob("*")) export_dir = str(children[0]) logger.info(f"load saved_model from {export_dir}") self.saved_model = tf.saved_model.load(export_dir=export_dir)
def main(unused_argv): del unused_argv # Load the label map. print(' - Loading the label map...') label_map_dict = {} if FLAGS.label_map_format == 'csv': with tf.gfile.Open(FLAGS.label_map_file, 'r') as csv_file: reader = csv.reader(csv_file, delimiter=':') for row in reader: if len(row) != 2: raise ValueError( 'Each row of the csv label map file must be in ' '`id:name` format.') id_index = int(row[0]) name = row[1] label_map_dict[id_index] = { 'id': id_index, 'name': name, } else: raise ValueError('Unsupported label map format: {}.'.format( FLAGS.label_mape_format)) params = config_factory.config_generator(FLAGS.model) if FLAGS.config_file: params = params_dict.override_params_dict(params, FLAGS.config_file, is_strict=True) params = params_dict.override_params_dict(params, FLAGS.params_override, is_strict=True) params.validate() params.lock() model = model_factory.model_generator(params) with tf.Graph().as_default(): image_input = tf.placeholder(shape=(), dtype=tf.string) image = tf.io.decode_image(image_input, channels=3) image.set_shape([None, None, 3]) image = input_utils.normalize_image(image) image_size = [FLAGS.image_size, FLAGS.image_size] image, image_info = input_utils.resize_and_crop_image( image, image_size, image_size, aug_scale_min=1.0, aug_scale_max=1.0) image.set_shape([image_size[0], image_size[1], 3]) # batching. images = tf.reshape(image, [1, image_size[0], image_size[1], 3]) images_info = tf.expand_dims(image_info, axis=0) # model inference outputs = model.build_outputs(images, {'image_info': images_info}, mode=mode_keys.PREDICT) outputs['detection_boxes'] = ( outputs['detection_boxes'] / tf.tile(images_info[:, 2:3, :], [1, 1, 2])) predictions = outputs # Create a saver in order to load the pre-trained checkpoint. saver = tf.train.Saver() image_with_detections_list = [] with tf.Session() as sess: print(' - Loading the checkpoint...') saver.restore(sess, FLAGS.checkpoint_path) image_files = tf.gfile.Glob(FLAGS.image_file_pattern) for i, image_file in enumerate(image_files): print(' - Processing image %d...' % i) with tf.gfile.GFile(image_file, 'rb') as f: image_bytes = f.read() image = Image.open(image_file) image = image.convert( 'RGB') # needed for images with 4 channels. width, height = image.size np_image = (np.array(image.getdata()).reshape( height, width, 3).astype(np.uint8)) predictions_np = sess.run(predictions, feed_dict={image_input: image_bytes}) num_detections = int(predictions_np['num_detections'][0]) np_boxes = predictions_np['detection_boxes'][ 0, :num_detections] np_scores = predictions_np['detection_scores'][ 0, :num_detections] np_classes = predictions_np['detection_classes'][ 0, :num_detections] np_classes = np_classes.astype(np.int32) np_masks = None if 'detection_masks' in predictions_np: instance_masks = predictions_np['detection_masks'][ 0, :num_detections] np_masks = mask_utils.paste_instance_masks( instance_masks, box_utils.yxyx_to_xywh(np_boxes), height, width) image_with_detections = ( visualization_utils. visualize_boxes_and_labels_on_image_array( np_image, np_boxes, np_classes, np_scores, label_map_dict, instance_masks=np_masks, use_normalized_coordinates=False, max_boxes_to_draw=FLAGS.max_boxes_to_draw, min_score_thresh=FLAGS.min_score_threshold)) image_with_detections_list.append(image_with_detections) print(' - Saving the outputs...') formatted_image_with_detections_list = [ Image.fromarray(image.astype(np.uint8)) for image in image_with_detections_list ] html_str = '<html>' image_strs = [] for formatted_image in formatted_image_with_detections_list: with io.BytesIO() as stream: formatted_image.save(stream, format='JPEG') data_uri = base64.b64encode(stream.getvalue()).decode('utf-8') image_strs.append( '<img src="data:image/jpeg;base64,{}", height=800>'.format( data_uri)) images_str = ' '.join(image_strs) html_str += images_str html_str += '</html>' with tf.gfile.GFile(FLAGS.output_html, 'w') as f: f.write(html_str)
def main(argv): del argv # Unused. params = factory.config_generator(FLAGS.model) if FLAGS.config_file: params = params_dict.override_params_dict(params, FLAGS.config_file, is_strict=True) params = params_dict.override_params_dict(params, FLAGS.params_override, is_strict=True) params.override( { 'platform': { 'eval_master': FLAGS.eval_master, 'tpu': FLAGS.tpu, 'tpu_zone': FLAGS.tpu_zone, 'gcp_project': FLAGS.gcp_project, }, 'tpu_job_name': FLAGS.tpu_job_name, 'use_tpu': FLAGS.use_tpu, 'model_dir': FLAGS.model_dir, 'train': { 'num_shards': FLAGS.num_cores, }, }, is_strict=False) # Only run spatial partitioning in training mode. if FLAGS.mode != 'train': params.train.input_partition_dims = None params.train.num_cores_per_replica = None params.validate() params.lock() pp = pprint.PrettyPrinter() params_str = pp.pformat(params.as_dict()) tf.logging.info('Model Parameters: {}'.format(params_str)) # Builds detection model on TPUs. model_fn = model_builder.ModelFn(params) executor = tpu_executor.TpuExecutor(model_fn, params) # Prepares input functions for train and eval. train_input_fn = input_reader.InputFn(params.train.train_file_pattern, params, mode=ModeKeys.TRAIN) eval_input_fn = input_reader.InputFn(params.eval.eval_file_pattern, params, mode=ModeKeys.PREDICT_WITH_GT) # Runs the model. if FLAGS.mode == 'train': save_config(params, params.model_dir) executor.train(train_input_fn, params.train.total_steps) if FLAGS.eval_after_training: executor.prepare_evaluation() executor.evaluate( eval_input_fn, params.eval.eval_samples // params.eval.eval_batch_size) elif FLAGS.mode == 'eval': def terminate_eval(): tf.logging.info( 'Terminating eval after %d seconds of no checkpoints' % params.eval.eval_timeout) return True executor.prepare_evaluation() # Runs evaluation when there's a new checkpoint. for ckpt in tf.contrib.training.checkpoints_iterator( params.model_dir, min_interval_secs=params.eval.min_eval_interval, timeout=params.eval.eval_timeout, timeout_fn=terminate_eval): # Terminates eval job when final checkpoint is reached. current_step = int(os.path.basename(ckpt).split('-')[1]) tf.logging.info('Starting to evaluate.') try: executor.evaluate( eval_input_fn, params.eval.eval_samples // params.eval.eval_batch_size, ckpt) if current_step >= params.train.total_steps: tf.logging.info( 'Evaluation finished after training step %d' % current_step) break except tf.errors.NotFoundError: # Since the coordinator is on a different job than the TPU worker, # sometimes the TPU worker does not finish initializing until long after # the CPU job tells it to start evaluating. In this case, the checkpoint # file could have been deleted already. tf.logging.info( 'Checkpoint %s no longer exists, skipping checkpoint' % ckpt) elif FLAGS.mode == 'train_and_eval': save_config(params, params.model_dir) executor.prepare_evaluation() num_cycles = int(params.train.total_steps / params.eval.num_steps_per_eval) for cycle in range(num_cycles): tf.logging.info('Start training cycle %d.' % cycle) current_cycle_last_train_step = ((cycle + 1) * params.eval.num_steps_per_eval) executor.train(train_input_fn, current_cycle_last_train_step) executor.evaluate( eval_input_fn, params.eval.eval_samples // params.eval.eval_batch_size) else: tf.logging.info('Mode not found.')
def main(argv): del argv # Unused. params = factory.config_generator(FLAGS.model) if FLAGS.config_file: params = params_dict.override_params_dict( params, FLAGS.config_file, is_strict=True) params = params_dict.override_params_dict( params, FLAGS.params_override, is_strict=True) if not FLAGS.use_tpu: params.override({ 'architecture': { 'use_bfloat16': False, }, 'batch_norm_activation': { 'use_sync_bn': False, }, }, is_strict=True) params.override({ 'platform': { 'eval_master': FLAGS.eval_master, 'tpu': FLAGS.tpu, 'tpu_zone': FLAGS.tpu_zone, 'gcp_project': FLAGS.gcp_project, }, 'tpu_job_name': FLAGS.tpu_job_name, 'use_tpu': FLAGS.use_tpu, 'model_dir': FLAGS.model_dir, 'train': { 'num_shards': FLAGS.num_cores, }, }, is_strict=False) # Only run spatial partitioning in training mode. if FLAGS.mode != 'train': params.train.input_partition_dims = None params.train.num_cores_per_replica = None params.validate() params.lock() pp = pprint.PrettyPrinter() params_str = pp.pformat(params.as_dict()) logging.info('Model Parameters: %s', params_str) # Builds detection model on TPUs. model_fn = model_builder.ModelFn(params) executor = tpu_executor.TpuExecutor(model_fn, params) # Prepares input functions for train and eval. train_input_fn = input_reader.InputFn( params.train.train_file_pattern, params, mode=ModeKeys.TRAIN, dataset_type=params.train.train_dataset_type) if params.eval.type == 'customized': eval_input_fn = input_reader.InputFn( params.eval.eval_file_pattern, params, mode=ModeKeys.EVAL, dataset_type=params.eval.eval_dataset_type) else: eval_input_fn = input_reader.InputFn( params.eval.eval_file_pattern, params, mode=ModeKeys.PREDICT_WITH_GT, dataset_type=params.eval.eval_dataset_type) # Runs the model. if FLAGS.mode == 'train': config_utils.save_config(params, params.model_dir) executor.train(train_input_fn, params.train.total_steps) if FLAGS.eval_after_training: executor.evaluate( eval_input_fn, params.eval.eval_samples // params.eval.eval_batch_size) elif FLAGS.mode == 'eval': def terminate_eval(): logging.info('Terminating eval after %d seconds of no checkpoints', params.eval.eval_timeout) return True # Runs evaluation when there's a new checkpoint. for ckpt in tf.train.checkpoints_iterator( params.model_dir, min_interval_secs=params.eval.min_eval_interval, timeout=params.eval.eval_timeout, timeout_fn=terminate_eval): # Terminates eval job when final checkpoint is reached. current_step = int(os.path.basename(ckpt).split('-')[1]) logging.info('Starting to evaluate.') try: executor.evaluate( eval_input_fn, params.eval.eval_samples // params.eval.eval_batch_size, ckpt) if current_step >= params.train.total_steps: logging.info('Evaluation finished after training step %d', current_step) break except tf.errors.NotFoundError: # Since the coordinator is on a different job than the TPU worker, # sometimes the TPU worker does not finish initializing until long after # the CPU job tells it to start evaluating. In this case, the checkpoint # file could have been deleted already. logging.info('Checkpoint %s no longer exists, skipping checkpoint', ckpt) elif FLAGS.mode == 'train_and_eval': config_utils.save_config(params, params.model_dir) num_cycles = int(params.train.total_steps / params.eval.num_steps_per_eval) for cycle in range(num_cycles): logging.info('Start training cycle %d.', cycle) current_cycle_last_train_step = ((cycle + 1) * params.eval.num_steps_per_eval) executor.train(train_input_fn, current_cycle_last_train_step) executor.evaluate( eval_input_fn, params.eval.eval_samples // params.eval.eval_batch_size) elif FLAGS.mode == 'predict': file_pattern = FLAGS.predict_file_pattern if not file_pattern: raise ValueError('"predict_file_pattern" parameter is required.') output_dir = FLAGS.predict_output_dir if not output_dir: raise ValueError('"predict_output_dir" parameter is required.') test_input_fn = input_reader.InputFn( file_pattern, params, mode=ModeKeys.PREDICT_WITH_GT, dataset_type=params.eval.eval_dataset_type) checkpoint_prefix = 'model.ckpt-' + FLAGS.predict_checkpoint_step checkpoint_path = os.path.join(FLAGS.model_dir, checkpoint_prefix) if not tf.train.checkpoint_exists(checkpoint_path): checkpoint_path = os.path.join(FLAGS.model_dir, 'best_checkpoints', checkpoint_prefix) if not tf.train.checkpoint_exists(checkpoint_path): raise ValueError('Checkpoint not found: %s/%s' % (FLAGS.model_dir, checkpoint_prefix)) executor.predict(test_input_fn, checkpoint_path, output_dir=output_dir) else: logging.info('Mode not found.')
def export(export_dir, checkpoint_path, model, config_file='', params_override='', use_tpu=False, batch_size=1, image_size=(1024, 1024), input_type='raw_image_tensor', input_name='input', output_image_info=True, output_normalized_coordinates=False, cast_num_detections_to_float=False, cast_detection_classes_to_float=False): """Exports the SavedModel.""" control_flow_util.enable_control_flow_v2() params = factory.config_generator(model) if config_file: params = params_dict.override_params_dict( params, config_file, is_strict=True) # Use `is_strict=False` to load params_override with run_time variables like # `train.num_shards`. params = params_dict.override_params_dict( params, params_override, is_strict=False) if not use_tpu: params.override({ 'architecture': { 'use_bfloat16': use_tpu, }, }, is_strict=True) if batch_size is None: params.override({ 'postprocess': { 'use_batched_nms': True, } }) params.validate() params.lock() model_params = dict( params.as_dict(), use_tpu=use_tpu, mode=tf.estimator.ModeKeys.PREDICT, transpose_input=False) tf.logging.info('model_params is:\n %s', model_params) if model in ['attribute_mask_rcnn']: model_fn = serving.serving_model_fn_builder( use_tpu, output_image_info, output_normalized_coordinates, cast_num_detections_to_float, cast_detection_classes_to_float) serving_input_receiver_fn = functools.partial( serving.serving_input_fn, batch_size=batch_size, desired_image_size=image_size, stride=(2 ** params.architecture.max_level), input_type=input_type, input_name=input_name) else: raise ValueError('The model type `{} is not supported.'.format(model)) print(' - Setting up TPUEstimator...') estimator = tf.estimator.tpu.TPUEstimator( model_fn=model_fn, model_dir=None, config=tf.estimator.tpu.RunConfig( tpu_config=tf.estimator.tpu.TPUConfig(iterations_per_loop=1), master='local', evaluation_master='local'), params=model_params, use_tpu=use_tpu, train_batch_size=batch_size, predict_batch_size=batch_size, export_to_tpu=use_tpu, export_to_cpu=True) print(' - Exporting the model...') dir_name = os.path.dirname(export_dir) if not tf.gfile.Exists(dir_name): tf.logging.info('Creating base dir: %s', dir_name) tf.gfile.MakeDirs(dir_name) export_path = estimator.export_saved_model( export_dir_base=dir_name, serving_input_receiver_fn=serving_input_receiver_fn, checkpoint_path=checkpoint_path) tf.logging.info( 'Exported SavedModel to %s, renaming to %s', export_path, export_dir) if tf.gfile.Exists(export_dir): tf.logging.info('Deleting existing SavedModel dir: %s', export_dir) tf.gfile.DeleteRecursively(export_dir) tf.gfile.Rename(export_path, export_dir)
def main(unused_argv): del unused_argv params = config_factory.config_generator(FLAGS.model) if FLAGS.config_file: params = params_dict.override_params_dict( params, FLAGS.config_file, is_strict=True) params = params_dict.override_params_dict( params, FLAGS.params_override, is_strict=True) # We currently only support batch_size = 1 to evaluate images one by one. # Override the `eval_batch_size` = 1 here. params.override({ 'eval': { 'eval_batch_size': 1, }, }) params.validate() params.lock() model = model_factory.model_generator(params) evaluator = evaluator_factory.evaluator_generator(params.eval) parse_fn = functools.partial(parse_single_example, params=params) with tf.Graph().as_default(): dataset = tf.data.Dataset.list_files( params.eval.eval_file_pattern, shuffle=False) dataset = dataset.apply( tf.data.experimental.parallel_interleave( lambda filename: tf.data.TFRecordDataset(filename).prefetch(1), cycle_length=32, sloppy=False)) dataset = dataset.map(parse_fn, num_parallel_calls=64) dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE) dataset = dataset.batch(1, drop_remainder=False) images, labels, groundtruths = dataset.make_one_shot_iterator().get_next() images.set_shape([ 1, params.retinanet_parser.output_size[0], params.retinanet_parser.output_size[1], 3]) # model inference outputs = model.build_outputs(images, labels, mode=mode_keys.PREDICT) predictions = outputs predictions.update({ 'source_id': groundtruths['source_id'], 'image_info': labels['image_info'], }) # Create a saver in order to load the pre-trained checkpoint. saver = tf.train.Saver() with tf.Session() as sess: saver.restore(sess, FLAGS.checkpoint_path) num_batches = params.eval.eval_samples // params.eval.eval_batch_size for i in range(num_batches): if i % 100 == 0: print('{}/{} batches...'.format(i, num_batches)) predictions_np, groundtruths_np = sess.run([predictions, groundtruths]) evaluator.update(predictions_np, groundtruths_np) if FLAGS.dump_predictions_only: print('Dumping the predction results...') evaluator.dump_predictions(FLAGS.predictions_path) print('Done!') else: print('Evaluating the prediction results...') metrics = evaluator.evaluate() print('Eval results: {}'.format(metrics))
def main(unused_argv): del unused_argv # Load the label map. print(' - Loading the label map...') label_map_dict = {} if FLAGS.label_map_format == 'csv': with tf.gfile.Open(FLAGS.label_map_file, 'r') as csv_file: reader = csv.reader(csv_file, delimiter=':') for row in reader: if len(row) != 2: raise ValueError( 'Each row of the csv label map file must be in ' '`id:name` format.') id_index = int(row[0]) name = row[1] label_map_dict[id_index] = { 'id': id_index, 'name': name, } else: raise ValueError('Unsupported label map format: {}.'.format( FLAGS.label_mape_format)) params = config_factory.config_generator(FLAGS.model) if FLAGS.config_file: params = params_dict.override_params_dict(params, FLAGS.config_file, is_strict=True) params = params_dict.override_params_dict(params, FLAGS.params_override, is_strict=True) params.override( { 'architecture': { 'use_bfloat16': False, # The inference runs on CPU/GPU. }, }, is_strict=True) params.validate() params.lock() model = model_factory.model_generator(params) with tf.Graph().as_default(): image_input = tf.placeholder(shape=(), dtype=tf.string) image = tf.io.decode_image(image_input, channels=3) image.set_shape([None, None, 3]) image = input_utils.normalize_image(image) image_size = [FLAGS.image_size, FLAGS.image_size] image, image_info = input_utils.resize_and_crop_image( image, image_size, image_size, aug_scale_min=1.0, aug_scale_max=1.0) image.set_shape([image_size[0], image_size[1], 3]) # batching. images = tf.reshape(image, [1, image_size[0], image_size[1], 3]) images_info = tf.expand_dims(image_info, axis=0) # model inference outputs = model.build_outputs(images, {'image_info': images_info}, mode=mode_keys.PREDICT) # outputs['detection_boxes'] = ( # outputs['detection_boxes'] / tf.tile(images_info[:, 2:3, :], [1, 1, 2])) predictions = outputs # Create a saver in order to load the pre-trained checkpoint. saver = tf.train.Saver() image_with_detections_list = [] with tf.Session() as sess: print(' - Loading the checkpoint...') saver.restore(sess, FLAGS.checkpoint_path) image_files = tf.gfile.Glob(FLAGS.image_file_pattern) for i, image_file in enumerate(image_files): print(' - Processing image %d...' % i) with tf.gfile.GFile(image_file, 'rb') as f: image_bytes = f.read() image = Image.open(image_file) image = image.convert( 'RGB') # needed for images with 4 channels. width, height = image.size np_image = (np.array(image.getdata()).reshape( height, width, 3).astype(np.uint8)) print(np_image.shape) predictions_np = sess.run(predictions, feed_dict={image_input: image_bytes}) logits = predictions_np['logits'][0] print(logits.shape) labels = np.argmax(logits.squeeze(), -1) print(labels.shape) print(labels) labels = np.array(Image.fromarray(labels.astype('uint8'))) print(labels.shape) plt.imshow(labels) plt.savefig(f"temp-{i}.png")
def main(argv): del argv # Unused. params = factory.config_generator(FLAGS.model) if FLAGS.config_file: params = params_dict.override_params_dict(params, FLAGS.config_file, is_strict=True) # Use `is_strict=False` to load params_override with run_time variables like # `train.num_shards`. params = params_dict.override_params_dict(params, FLAGS.params_override, is_strict=False) params.validate() params.lock() image_size = [int(x) for x in FLAGS.input_image_size.split(',')] g = tf.Graph() with g.as_default(): # Build the input. _, features = inputs.build_serving_input( input_type=FLAGS.input_type, batch_size=FLAGS.batch_size, desired_image_size=image_size, stride=(2**params.anchor.max_level)) # Build the model. print(' - Building the graph...') if FLAGS.model in ['retinanet', 'mask_rcnn', 'shapemask']: graph_fn = detection.serving_model_graph_builder( FLAGS.output_image_info, FLAGS.output_normalized_coordinates, FLAGS.cast_num_detections_to_float) else: raise ValueError('The model type `{}` is not supported.'.format( FLAGS.model)) predictions = graph_fn(features, params) # Add a saver for checkpoint loading. tf.train.Saver() inference_graph_def = g.as_graph_def() optimized_graph_def = inference_graph_def if FLAGS.optimize_graph: print(' - Optimizing the graph...') # Trim the unused nodes in the graph. output_nodes = [ output_node.op.name for output_node in predictions.values() ] # TODO(pengchong): Consider to use `strip_unused_lib.strip_unused` and/or # `optimize_for_inference_lib.optimize_for_inference` to trim the graph. # Use `optimize_for_inference` if we decide to export the frozen graph # (graph + checkpoint) and want explictily fold in batchnorm variables. optimized_graph_def = graph_util.remove_training_nodes( optimized_graph_def, output_nodes) print(' - Saving the graph...') tf.train.write_graph(optimized_graph_def, FLAGS.export_dir, 'inference_graph.pbtxt') print(' - Done!')
def main(argv): del argv # Unused. params = factory.config_generator(FLAGS.model) if FLAGS.config_file: params = params_dict.override_params_dict(params, FLAGS.config_file, is_strict=True) # Use `is_strict=False` to load params_override with run_time variables like # `train.num_shards`. params = params_dict.override_params_dict(params, FLAGS.params_override, is_strict=False) params.validate() params.lock() model_params = dict(params.as_dict(), use_tpu=FLAGS.use_tpu, mode=tf.estimator.ModeKeys.PREDICT, transpose_input=False) tf.logging.info('model_params is:\n %s', model_params) image_size = [int(x) for x in FLAGS.input_image_size.split(',')] if FLAGS.model == 'retinanet': model_fn = serving.serving_model_fn_builder( FLAGS.use_tpu, FLAGS.output_image_info, FLAGS.output_normalized_coordinates, FLAGS.cast_num_detections_to_float) serving_input_receiver_fn = functools.partial( serving.serving_input_fn, batch_size=FLAGS.batch_size, desired_image_size=image_size, stride=(2**params.anchor.max_level), input_type=FLAGS.input_type, input_name=FLAGS.input_name) else: raise ValueError('Model %s is not supported.' % params.type) print(' - Setting up TPUEstimator...') estimator = tf.estimator.tpu.TPUEstimator( model_fn=model_fn, model_dir=None, config=tf.estimator.tpu.RunConfig( tpu_config=tf.estimator.tpu.TPUConfig(iterations_per_loop=1), master='local', evaluation_master='local'), params=model_params, use_tpu=FLAGS.use_tpu, train_batch_size=FLAGS.batch_size, predict_batch_size=FLAGS.batch_size, export_to_tpu=FLAGS.use_tpu, export_to_cpu=True) print(' - Exporting the model...') dir_name = os.path.dirname(FLAGS.export_dir) if not tf.gfile.Exists(dir_name): tf.logging.info('Creating base dir: %s', dir_name) tf.gfile.MakeDirs(dir_name) export_path = estimator.export_saved_model( export_dir_base=dir_name, serving_input_receiver_fn=serving_input_receiver_fn, checkpoint_path=FLAGS.checkpoint_path) tf.logging.info('Exported SavedModel to %s, renaming to %s', export_path, FLAGS.export_dir) if tf.gfile.Exists(FLAGS.export_dir): tf.logging.info('Deleting existing SavedModel dir: %s', FLAGS.export_dir) tf.gfile.DeleteRecursively(FLAGS.export_dir) tf.gfile.Rename(export_path, FLAGS.export_dir)
def initiate(): # Load the label map. print(' - Loading the label map...') label_map_dict = {} if 'csv' == 'csv': with tf.gfile.Open('dataset/fashionpedia_label_map.csv', 'r') as csv_file: reader = csv.reader(csv_file, delimiter=':') for row in reader: if len(row) != 2: raise ValueError( 'Each row of the csv label map file must be in ' '`id:name` format.') id_index = int(row[0]) name = row[1] label_map_dict[id_index] = { 'id': id_index, 'name': name, } else: raise ValueError('Unsupported label map format: {}.'.format('csv')) params = config_factory.config_generator('attribute_mask_rcnn') if 'configs/yaml/spinenet49_amrcnn.yaml': params = params_dict.override_params_dict( params, 'configs/yaml/spinenet49_amrcnn.yaml', is_strict=True) params = params_dict.override_params_dict(params, '', is_strict=True) params.override( { 'architecture': { 'use_bfloat16': False, # The inference runs on CPU/GPU. }, }, is_strict=True) params.validate() params.lock() model = model_factory.model_generator(params) with tf.Graph().as_default(): image_input = tf.placeholder(shape=(), dtype=tf.string) image = tf.io.decode_image(image_input, channels=3) image.set_shape([None, None, 3]) image = input_utils.normalize_image(image) image_size = [640, 640] image, image_info = input_utils.resize_and_crop_image( image, image_size, image_size, aug_scale_min=1.0, aug_scale_max=1.0) image.set_shape([image_size[0], image_size[1], 3]) # batching. images = tf.reshape(image, [1, image_size[0], image_size[1], 3]) images_info = tf.expand_dims(image_info, axis=0) # model inference outputs = model.build_outputs(images, {'image_info': images_info}, mode=mode_keys.PREDICT) outputs['detection_boxes'] = ( outputs['detection_boxes'] / tf.tile(images_info[:, 2:3, :], [1, 1, 2])) predictions = outputs # Create a saver in order to load the pre-trained checkpoint. saver = tf.train.Saver() sess = tf.Session() print(' - Loading the checkpoint...') saver.restore(sess, 'fashionpedia-spinenet-49/model.ckpt') print(' - Checkpoint Loaded...') return sess, predictions, image_input
def main(argv): del argv # Unused. params = factory.config_generator(FLAGS.model) if FLAGS.config_file: params = params_dict.override_params_dict(params, FLAGS.config_file, is_strict=True) params = params_dict.override_params_dict(params, FLAGS.params_override, is_strict=True) params.override({ 'use_tpu': FLAGS.use_tpu, 'model_dir': FLAGS.model_dir, }, is_strict=True) if not FLAGS.use_tpu: params.override( { 'architecture': { 'use_bfloat16': False, }, 'batch_norm_activation': { 'use_sync_bn': False, }, }, is_strict=True) # Only run spatial partitioning in training mode. if FLAGS.mode != 'train': params.train.input_partition_dims = None params.train.num_cores_per_replica = None params_to_save = params_dict.ParamsDict(params) params.override( { 'platform': { 'eval_master': FLAGS.eval_master, 'tpu': FLAGS.tpu, 'tpu_zone': FLAGS.tpu_zone, 'gcp_project': FLAGS.gcp_project, }, 'tpu_job_name': FLAGS.tpu_job_name, 'train': { 'num_shards': FLAGS.num_cores, }, }, is_strict=False) params.validate() params.lock() pp = pprint.PrettyPrinter() params_str = pp.pformat(params.as_dict()) logging.info('Model Parameters: %s', params_str) # Builds detection model on TPUs. model_fn = model_builder.ModelFn(params) executor = tpu_executor.TpuExecutor(model_fn, params) # Prepares input functions for train and eval. train_input_fn = input_reader.InputFn( params.train.train_file_pattern, params, mode=ModeKeys.TRAIN, dataset_type=params.train.train_dataset_type) if params.eval.type == 'customized': eval_input_fn = input_reader.InputFn( params.eval.eval_file_pattern, params, mode=ModeKeys.EVAL, dataset_type=params.eval.eval_dataset_type) else: eval_input_fn = input_reader.InputFn( params.eval.eval_file_pattern, params, mode=ModeKeys.PREDICT_WITH_GT, dataset_type=params.eval.eval_dataset_type) if params.eval.eval_samples: eval_times = params.eval.eval_samples // params.eval.eval_batch_size else: eval_times = None # Runs the model. if FLAGS.mode == 'train': config_utils.save_config(params_to_save, params.model_dir) executor.train(train_input_fn, params.train.total_steps) if FLAGS.eval_after_training: executor.evaluate(eval_input_fn, eval_times) elif FLAGS.mode == 'eval': def terminate_eval(): logging.info('Terminating eval after %d seconds of no checkpoints', params.eval.eval_timeout) return True # Runs evaluation when there's a new checkpoint. for ckpt in tf.train.checkpoints_iterator( params.model_dir, min_interval_secs=params.eval.min_eval_interval, timeout=params.eval.eval_timeout, timeout_fn=terminate_eval): # Terminates eval job when final checkpoint is reached. current_step = int( six.ensure_str(os.path.basename(ckpt)).split('-')[1]) logging.info('Starting to evaluate.') try: executor.evaluate(eval_input_fn, eval_times, ckpt) if current_step >= params.train.total_steps: logging.info('Evaluation finished after training step %d', current_step) break except tf.errors.NotFoundError as e: logging.info( 'Erorr occurred during evaluation: NotFoundError: %s', e) elif FLAGS.mode == 'train_and_eval': config_utils.save_config(params_to_save, params.model_dir) num_cycles = int(params.train.total_steps / params.eval.num_steps_per_eval) for cycle in range(num_cycles): logging.info('Start training cycle %d.', cycle) current_cycle_last_train_step = ((cycle + 1) * params.eval.num_steps_per_eval) executor.train(train_input_fn, current_cycle_last_train_step) executor.evaluate(eval_input_fn, eval_times) else: logging.info('Mode not found.')
def main(argv): del argv # Unused. params = factory.config_generator(FLAGS.model) if FLAGS.config_file: params = params_dict.override_params_dict( params, FLAGS.config_file, is_strict=True) params = params_dict.override_params_dict( params, FLAGS.params_override, is_strict=True) params.train.input_partition_dims = None params.train.num_cores_per_replica = None params.architecture.use_bfloat16 = False # params.maskrcnn_parser.use_autoaugment = False params.validate() params.lock() # Prepares input functions for train and eval. train_input_fn = input_reader.InputFnTest( params.train.train_file_pattern, params, mode=ModeKeys.TRAIN, dataset_type=params.train.train_dataset_type) batch_size = 1 dataset = train_input_fn({'batch_size': batch_size}) category_index = {} for i in range(50): category_index[i] = { 'name': 'test_%d' % i, } for i, (image_batch, labels_batch) in enumerate(dataset.take(10)): image_batch = tf.transpose(image_batch, [3, 0, 1, 2]) image_batch = tf.map_fn(denormalize_image, image_batch, dtype=tf.uint8, back_prop=False) image_shape = tf.shape(image_batch)[1:3] masks_batch = [] for image, bboxes, masks in zip(image_batch, labels_batch['gt_boxes'], labels_batch['gt_masks']): # extract masks bboxes = tf.numpy_function(box_utils.yxyx_to_xywh, [bboxes], tf.float32) binary_masks = tf.numpy_function(mask_utils.paste_instance_masks, [masks, bboxes, image_shape[0], image_shape[1]], tf.uint8) masks_batch.append(binary_masks) masks_batch = tf.stack(masks_batch, axis=0) scores_mask = tf.cast(tf.greater(labels_batch['gt_classes'], -1), tf.float32) scores = tf.ones_like(labels_batch['gt_classes'], dtype=tf.float32) * scores_mask images = draw_bounding_boxes_on_image_tensors(image_batch, labels_batch['gt_boxes'], labels_batch['gt_classes'], scores, category_index, instance_masks=masks_batch, use_normalized_coordinates=False) for j, image in enumerate(images): image_bytes = tf.io.encode_jpeg(image) tf.io.write_file(root_dir('data/visualizations/aug_%d.jpg' % (i * batch_size + j)), image_bytes)