def test_get_configs_from_pipeline_file(self): """Test that proto configs can be read from pipeline config file.""" pipeline_config_path = os.path.join(self.get_temp_dir(), "pipeline.config") pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() pipeline_config.model.faster_rcnn.num_classes = 10 pipeline_config.train_config.batch_size = 32 pipeline_config.train_input_reader.label_map_path = "path/to/label_map" pipeline_config.eval_config.num_examples = 20 pipeline_config.eval_input_reader.queue_capacity = 100 _write_config(pipeline_config, pipeline_config_path) configs = config_util.get_configs_from_pipeline_file( pipeline_config_path) self.assertProtoEquals(pipeline_config.model, configs["model"]) self.assertProtoEquals(pipeline_config.train_config, configs["train_config"]) self.assertProtoEquals(pipeline_config.train_input_reader, configs["train_input_config"]) self.assertProtoEquals(pipeline_config.eval_config, configs["eval_config"]) self.assertProtoEquals(pipeline_config.eval_input_reader, configs["eval_input_config"])
def testNewFocalLossParameters(self): """Tests that the loss weight ratio is updated appropriately.""" original_alpha = 1.0 original_gamma = 1.0 new_alpha = 0.3 new_gamma = 2.0 hparams = tf.contrib.training.HParams(focal_loss_alpha=new_alpha, focal_loss_gamma=new_gamma) pipeline_config_path = os.path.join(self.get_temp_dir(), "pipeline.config") pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() classification_loss = pipeline_config.model.ssd.loss.classification_loss classification_loss.weighted_sigmoid_focal.alpha = original_alpha classification_loss.weighted_sigmoid_focal.gamma = original_gamma _write_config(pipeline_config, pipeline_config_path) configs = config_util.get_configs_from_pipeline_file( pipeline_config_path) configs = config_util.merge_external_params_with_configs( configs, hparams) classification_loss = configs["model"].ssd.loss.classification_loss self.assertAlmostEqual( new_alpha, classification_loss.weighted_sigmoid_focal.alpha) self.assertAlmostEqual( new_gamma, classification_loss.weighted_sigmoid_focal.gamma)
def main(unused_argv): assert FLAGS.checkpoint_dir, '`checkpoint_dir` is missing.' assert FLAGS.eval_dir, '`eval_dir` is missing.' tf.gfile.MakeDirs(FLAGS.eval_dir) if FLAGS.pipeline_config_path: configs = config_util.get_configs_from_pipeline_file( FLAGS.pipeline_config_path) tf.gfile.Copy( FLAGS.pipeline_config_path, os.path.join(FLAGS.eval_dir, 'pipeline.config'), overwrite=True) 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) for name, config in [('model.config', FLAGS.model_config_path), ('eval.config', FLAGS.eval_config_path), ('input.config', FLAGS.input_config_path)]: tf.gfile.Copy(config, os.path.join(FLAGS.eval_dir, name), overwrite=True) model_config = configs['model'] eval_config = configs['eval_config'] input_config = configs['eval_input_config'] if FLAGS.eval_training_data: input_config = configs['train_input_config'] model_fn = functools.partial( model_builder.build, model_config=model_config, is_training=False) def get_next(config): return dataset_builder.make_initializable_iterator( dataset_builder.build(config)).get_next() create_input_dict_fn = functools.partial(get_next, input_config) categories = label_map_util.create_categories_from_labelmap( input_config.label_map_path) if FLAGS.run_once: eval_config.max_evals = 1 graph_rewriter_fn = None if 'graph_rewriter_config' in configs: graph_rewriter_fn = graph_rewriter_builder.build( configs['graph_rewriter_config'], is_training=False) evaluator.evaluate( create_input_dict_fn, model_fn, eval_config, categories, FLAGS.checkpoint_dir, FLAGS.eval_dir, graph_hook_fn=graph_rewriter_fn)
def testGetNumberOfClasses(self): """Tests that number of classes can be retrieved.""" pipeline_config_path = os.path.join(self.get_temp_dir(), "pipeline.config") pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() pipeline_config.model.faster_rcnn.num_classes = 20 _write_config(pipeline_config, pipeline_config_path) configs = config_util.get_configs_from_pipeline_file( pipeline_config_path) number_of_classes = config_util.get_number_of_classes(configs["model"]) self.assertEqual(20, number_of_classes)
def _get_configs_for_model(model_name): """Returns configurations for model.""" filename = model_test_util.GetPipelineConfigPath(model_name) data_path = _get_data_path() label_map_path = _get_labelmap_path() configs = config_util.get_configs_from_pipeline_file(filename) configs = config_util.merge_external_params_with_configs( configs, train_input_path=data_path, eval_input_path=data_path, label_map_path=label_map_path) return configs
def test_save_pipeline_config(self): """Tests that the pipeline config is properly saved to disk.""" pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() pipeline_config.model.faster_rcnn.num_classes = 10 pipeline_config.train_config.batch_size = 32 pipeline_config.train_input_reader.label_map_path = "path/to/label_map" pipeline_config.eval_config.num_examples = 20 pipeline_config.eval_input_reader.queue_capacity = 100 config_util.save_pipeline_config(pipeline_config, self.get_temp_dir()) configs = config_util.get_configs_from_pipeline_file( os.path.join(self.get_temp_dir(), "pipeline.config")) pipeline_config_reconstructed = ( config_util.create_pipeline_proto_from_configs(configs)) self.assertEqual(pipeline_config, pipeline_config_reconstructed)
def testNewBatchSizeWithClipping(self): """Tests that batch size is clipped to 1 from below.""" original_batch_size = 2 hparams = tf.contrib.training.HParams(batch_size=0.5) pipeline_config_path = os.path.join(self.get_temp_dir(), "pipeline.config") pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() pipeline_config.train_config.batch_size = original_batch_size _write_config(pipeline_config, pipeline_config_path) configs = config_util.get_configs_from_pipeline_file( pipeline_config_path) configs = config_util.merge_external_params_with_configs( configs, hparams) new_batch_size = configs["train_config"].batch_size self.assertEqual(1, new_batch_size) # Clipped to 1.0.
def testNewBatchSize(self): """Tests that batch size is updated appropriately.""" original_batch_size = 2 hparams = tf.contrib.training.HParams(batch_size=16) pipeline_config_path = os.path.join(self.get_temp_dir(), "pipeline.config") pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() pipeline_config.train_config.batch_size = original_batch_size _write_config(pipeline_config, pipeline_config_path) configs = config_util.get_configs_from_pipeline_file( pipeline_config_path) configs = config_util.merge_external_params_with_configs( configs, hparams) new_batch_size = configs["train_config"].batch_size self.assertEqual(16, new_batch_size)
def test_create_pipeline_proto_from_configs(self): """Tests that proto can be reconstructed from configs dictionary.""" pipeline_config_path = os.path.join(self.get_temp_dir(), "pipeline.config") pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() pipeline_config.model.faster_rcnn.num_classes = 10 pipeline_config.train_config.batch_size = 32 pipeline_config.train_input_reader.label_map_path = "path/to/label_map" pipeline_config.eval_config.num_examples = 20 pipeline_config.eval_input_reader.queue_capacity = 100 _write_config(pipeline_config, pipeline_config_path) configs = config_util.get_configs_from_pipeline_file( pipeline_config_path) pipeline_config_reconstructed = ( config_util.create_pipeline_proto_from_configs(configs)) self.assertEqual(pipeline_config, pipeline_config_reconstructed)
def _get_configs_for_model(model_name): """Returns configurations for model.""" # TODO: Make sure these tests work fine outside google3. fname = os.path.join( FLAGS.test_srcdir, ('google3/third_party/tensorflow_models/' 'object_detection/samples/configs/' + model_name + '.config')) label_map_path = os.path.join(FLAGS.test_srcdir, ('google3/third_party/tensorflow_models/' 'object_detection/data/pet_label_map.pbtxt')) data_path = os.path.join(FLAGS.test_srcdir, ('google3/third_party/tensorflow_models/' 'object_detection/test_data/pets_examples.record')) configs = config_util.get_configs_from_pipeline_file(fname) return config_util.merge_external_params_with_configs( configs, train_input_path=data_path, eval_input_path=data_path, label_map_path=label_map_path)
def testNewMomentumOptimizerValue(self): """Tests that new momentum value is updated appropriately.""" original_momentum_value = 0.4 hparams = tf.contrib.training.HParams(momentum_optimizer_value=1.1) pipeline_config_path = os.path.join(self.get_temp_dir(), "pipeline.config") pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() optimizer_config = pipeline_config.train_config.optimizer.rms_prop_optimizer optimizer_config.momentum_optimizer_value = original_momentum_value _write_config(pipeline_config, pipeline_config_path) configs = config_util.get_configs_from_pipeline_file( pipeline_config_path) configs = config_util.merge_external_params_with_configs( configs, hparams) optimizer_config = configs["train_config"].optimizer.rms_prop_optimizer new_momentum_value = optimizer_config.momentum_optimizer_value self.assertAlmostEqual(1.0, new_momentum_value) # Clipped to 1.0.
def testNewTrainInputPathList(self): """Tests that train input path can be overwritten with multiple files.""" original_train_path = ["path/to/data"] new_train_path = ["another/path/to/data", "yet/another/path/to/data"] pipeline_config_path = os.path.join(self.get_temp_dir(), "pipeline.config") pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() reader_config = pipeline_config.train_input_reader.tf_record_input_reader reader_config.input_path.extend(original_train_path) _write_config(pipeline_config, pipeline_config_path) configs = config_util.get_configs_from_pipeline_file( pipeline_config_path) configs = config_util.merge_external_params_with_configs( configs, train_input_path=new_train_path) reader_config = configs["train_input_config"].tf_record_input_reader final_path = reader_config.input_path self.assertEqual(new_train_path, final_path)
def testNewMaskType(self): """Tests that mask type can be overwritten in input readers.""" original_mask_type = input_reader_pb2.NUMERICAL_MASKS new_mask_type = input_reader_pb2.PNG_MASKS pipeline_config_path = os.path.join(self.get_temp_dir(), "pipeline.config") pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() train_input_reader = pipeline_config.train_input_reader train_input_reader.mask_type = original_mask_type eval_input_reader = pipeline_config.eval_input_reader eval_input_reader.mask_type = original_mask_type _write_config(pipeline_config, pipeline_config_path) configs = config_util.get_configs_from_pipeline_file( pipeline_config_path) configs = config_util.merge_external_params_with_configs( configs, mask_type=new_mask_type) self.assertEqual(new_mask_type, configs["train_input_config"].mask_type) self.assertEqual(new_mask_type, configs["eval_input_config"].mask_type)
def testNewLabelMapPath(self): """Tests that label map path can be overwritten in input readers.""" original_label_map_path = "path/to/original/label_map" new_label_map_path = "path//to/new/label_map" pipeline_config_path = os.path.join(self.get_temp_dir(), "pipeline.config") pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() train_input_reader = pipeline_config.train_input_reader train_input_reader.label_map_path = original_label_map_path eval_input_reader = pipeline_config.eval_input_reader eval_input_reader.label_map_path = original_label_map_path _write_config(pipeline_config, pipeline_config_path) configs = config_util.get_configs_from_pipeline_file( pipeline_config_path) configs = config_util.merge_external_params_with_configs( configs, label_map_path=new_label_map_path) self.assertEqual(new_label_map_path, configs["train_input_config"].label_map_path) self.assertEqual(new_label_map_path, configs["eval_input_config"].label_map_path)
def testNewClassificationLocalizationWeightRatio(self): """Tests that the loss weight ratio is updated appropriately.""" original_localization_weight = 0.1 original_classification_weight = 0.2 new_weight_ratio = 5.0 hparams = tf.contrib.training.HParams( classification_localization_weight_ratio=new_weight_ratio) pipeline_config_path = os.path.join(self.get_temp_dir(), "pipeline.config") pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() pipeline_config.model.ssd.loss.localization_weight = ( original_localization_weight) pipeline_config.model.ssd.loss.classification_weight = ( original_classification_weight) _write_config(pipeline_config, pipeline_config_path) configs = config_util.get_configs_from_pipeline_file( pipeline_config_path) configs = config_util.merge_external_params_with_configs( configs, hparams) loss = configs["model"].ssd.loss self.assertAlmostEqual(1.0, loss.localization_weight) self.assertAlmostEqual(new_weight_ratio, loss.classification_weight)
def testMergingKeywordArguments(self): """Tests that keyword arguments get merged as do hyperparameters.""" original_num_train_steps = 100 original_num_eval_steps = 5 desired_num_train_steps = 10 desired_num_eval_steps = 1 pipeline_config_path = os.path.join(self.get_temp_dir(), "pipeline.config") pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() pipeline_config.train_config.num_steps = original_num_train_steps pipeline_config.eval_config.num_examples = original_num_eval_steps _write_config(pipeline_config, pipeline_config_path) configs = config_util.get_configs_from_pipeline_file( pipeline_config_path) configs = config_util.merge_external_params_with_configs( configs, train_steps=desired_num_train_steps, eval_steps=desired_num_eval_steps) train_steps = configs["train_config"].num_steps eval_steps = configs["eval_config"].num_examples self.assertEqual(desired_num_train_steps, train_steps) self.assertEqual(desired_num_eval_steps, eval_steps)
def run_inference_for_single_image(image, graph, evaluator): with graph.as_default(): with tf.Session() as sess: configs = config_util.get_configs_from_pipeline_file( PIPELINE_CONFIG_PATH) def get_next(config): return dataset_util.make_initializable_iterator( dataset_builder.build(config)).get_next() create_input_dict_fn = functools.partial( get_next, configs['eval_input_config']) eval_config = configs['eval_config'] model_fn = functools.partial(model_builder.build, model_config=configs['model'], is_training=False) model = model_fn() tensor_dict = evtor._extract_prediction_tensors( model=model, create_input_dict_fn=create_input_dict_fn, ignore_groundtruth=eval_config.ignore_groundtruth) # Get handles to input and output tensors ops = tf.get_default_graph().get_operations() all_tensor_names = { output.name for op in ops for output in op.outputs } #tensor_dict = {} # for key in [ # 'num_detections', 'detection_boxes', 'detection_scores', # 'detection_classes', 'detection_masks' # ]: # tensor_name = key + ':0' # if tensor_name in all_tensor_names: # tensor_dict[key] = tf.get_default_graph().get_tensor_by_name( # tensor_name) if 'detection_masks' in tensor_dict: # The following processing is only for single image detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0]) detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0]) # Reframe is required to translate mask from box coordinates to image coordinates and fit the image size. real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32) detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1]) detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1]) detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks( detection_masks, detection_boxes, image.shape[0], image.shape[1]) detection_masks_reframed = tf.cast( tf.greater(detection_masks_reframed, 0.5), tf.uint8) # Follow the convention by adding back the batch dimension tensor_dict['detection_masks'] = tf.expand_dims( detection_masks_reframed, 0) image_tensor = tf.get_default_graph().get_tensor_by_name( 'image_tensor:0') # Run inference output_dict = sess.run( tensor_dict, feed_dict={image_tensor: np.expand_dims(image, 0)}) # all outputs are float32 numpy arrays, so convert types as appropriate output_dict['num_detections'] = int( output_dict['num_detections'][0]) output_dict['detection_classes'] = output_dict[ 'detection_classes'][0].astype(np.uint8) output_dict['detection_boxes'] = output_dict['detection_boxes'][0] output_dict['detection_scores'] = output_dict['detection_scores'][ 0] if 'detection_masks' in output_dict: output_dict['detection_masks'] = output_dict[ 'detection_masks'][0] evaluator.add_single_ground_truth_image_info( image_id=1, groundtruth_dict=output_dict) evaluator.add_single_detected_image_info( image_id=1, detections_dict=output_dict) metrics = evaluator.evaluate() evaluator.clear() return output_dict
def main(unused_argv): assert FLAGS.checkpoint_dir, '`checkpoint_dir` is missing.' assert FLAGS.eval_dir, '`eval_dir` is missing.' tf.gfile.MakeDirs(FLAGS.eval_dir) configs = config_util.get_configs_from_pipeline_file( FLAGS.pipeline_config_path) tf.gfile.Copy(FLAGS.pipeline_config_path, os.path.join(FLAGS.eval_dir, 'pipeline.config'), overwrite=True) model_config = configs['model'] input_config = configs['eval_input_config'] model_fn = functools.partial(model_builder.build, model_config=model_config, is_training=False) label_map = label_map = label_map_util.load_labelmap( '../input/oidv4label/oid_v4_label_map.pbtxt') 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, False) model = model_fn() image_raw_data = tf.placeholder(tf.string, None) img_data_jpg = tf.image.decode_jpeg(image_raw_data) input_img = tf.image.convert_image_dtype(img_data_jpg, dtype=tf.uint8) input_img.set_shape([None, None, 3]) original_image = tf.expand_dims(input_img, 0) preprocessed_image, true_image_shapes = model.preprocess( tf.to_float(original_image)) #preprocessed_image = tf.reshape(preprocessed_image, [-1, h, w, 3]) prediction_dict = model.predict(preprocessed_image, true_image_shapes) detections = model.postprocess(prediction_dict, true_image_shapes) label_id_offset = 1 output_dict = {} detection_fields = fields.DetectionResultFields detection_boxes = detections[detection_fields.detection_boxes][0] detection_scores = detections[detection_fields.detection_scores][0] detection_classes = ( tf.to_int64(detections[detection_fields.detection_classes][0]) + label_id_offset) num_detections = tf.to_int32( detections[detection_fields.num_detections][0]) detection_boxes = tf.slice(detection_boxes, begin=[0, 0], size=[num_detections, -1]) detection_classes = tf.slice(detection_classes, begin=[0], size=[num_detections]) detection_scores = tf.slice(detection_scores, begin=[0], size=[num_detections]) output_dict[detection_fields.detection_boxes] = detection_boxes output_dict[detection_fields.detection_classes] = detection_classes output_dict[detection_fields.detection_scores] = detection_scores variables_to_restore = tf.global_variables() global_step = tf.train.get_or_create_global_step() variables_to_restore.append(global_step) saver = tf.train.Saver(variables_to_restore) sess = tf.Session('', graph=tf.get_default_graph()) sess.run(tf.global_variables_initializer()) sess.run(tf.local_variables_initializer()) sess.run(tf.tables_initializer()) #latest_checkpoint = tf.train.latest_checkpoint(FLAGS.checkpoint_dir) latest_checkpoint = FLAGS.checkpoint_dir + '/model.ckpt' print("latest_checkpoint:" + latest_checkpoint) saver.restore(sess, latest_checkpoint) with open(FLAGS.eval_dir + '/test.csv', 'w') as f: re_str = 'ImageId,PredictionString\n' imgs = os.listdir(FLAGS.imgs_dir) img_count = 0 for img in imgs: img_count = img_count + 1 print(img_count) image_raw_data_jpg = tf.gfile.FastGFile(FLAGS.imgs_dir + '/' + img, 'rb').read() result_dict = sess.run( output_dict, feed_dict={image_raw_data: image_raw_data_jpg}) re_str = re_str + img.split(".")[0] re_str = re_str + ',' detected_box = result_dict[detection_fields.detection_boxes] for i in range(len(detected_box)): #print(result_dict[detection_fields.detection_classes][i]-1) #print(categories[0]) re_str = re_str + categories[result_dict[ detection_fields.detection_classes][i] - 1]['name'] + ' ' re_str = re_str + str( result_dict[detection_fields.detection_scores][i]) + ' ' re_str = re_str + str(detected_box[i][1]) + ' ' re_str = re_str + str(detected_box[i][0]) + ' ' re_str = re_str + str(detected_box[i][3]) + ' ' re_str = re_str + str(detected_box[i][2]) + ' ' f.write(re_str + '\n') re_str = '' sess.close()
def populate_experiment(run_config, hparams, pipeline_config_path, train_steps=None, eval_steps=None, model_fn_creator=create_model_fn, **kwargs): """Populates an `Experiment` object. Args: run_config: A `RunConfig`. hparams: A `HParams`. pipeline_config_path: A path to a pipeline config file. train_steps: Number of training steps. If None, the number of training steps is set from the `TrainConfig` proto. eval_steps: Number of evaluation steps per evaluation cycle. If None, the number of evaluation steps is set from the `EvalConfig` proto. model_fn_creator: A function that creates a `model_fn` for `Estimator`. Follows the signature: * Args: * `detection_model_fn`: Function that returns `DetectionModel` instance. * `configs`: Dictionary of pipeline config objects. * `hparams`: `HParams` object. * Returns: `model_fn` for `Estimator`. **kwargs: Additional keyword arguments for configuration override. Returns: An `Experiment` that defines all aspects of training, evaluation, and export. """ configs = config_util.get_configs_from_pipeline_file(pipeline_config_path) configs = config_util.merge_external_params_with_configs( configs, hparams, train_steps=train_steps, eval_steps=eval_steps, **kwargs) model_config = configs['model'] train_config = configs['train_config'] train_input_config = configs['train_input_config'] eval_config = configs['eval_config'] eval_input_config = configs['eval_input_config'] if train_steps is None: train_steps = train_config.num_steps if train_config.num_steps else None if eval_steps is None: eval_steps = eval_config.num_examples if eval_config.num_examples else None detection_model_fn = functools.partial( model_builder.build, model_config=model_config) # Create the input functions for TRAIN/EVAL. train_input_fn = inputs.create_train_input_fn( train_config=train_config, train_input_config=train_input_config, model_config=model_config) eval_input_fn = inputs.create_eval_input_fn( eval_config=eval_config, eval_input_config=eval_input_config, model_config=model_config) export_strategies = [ tf.contrib.learn.utils.saved_model_export_utils.make_export_strategy( serving_input_fn=inputs.create_predict_input_fn( model_config=model_config)) ] estimator = tf.estimator.Estimator( model_fn=model_fn_creator(detection_model_fn, configs, hparams), config=run_config) if run_config.is_chief: # Store the final pipeline config for traceability. pipeline_config_final = config_util.create_pipeline_proto_from_configs( configs) pipeline_config_final_path = os.path.join(estimator.model_dir, 'pipeline.config') config_text = text_format.MessageToString(pipeline_config_final) with tf.gfile.Open(pipeline_config_final_path, 'wb') as f: tf.logging.info('Writing as-run pipeline config file to %s', pipeline_config_final_path) f.write(config_text) return tf.contrib.learn.Experiment( estimator=estimator, train_input_fn=train_input_fn, eval_input_fn=eval_input_fn, train_steps=train_steps, eval_steps=eval_steps, export_strategies=export_strategies, eval_delay_secs=120,)
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'] train_config = configs['train_config'] input_config = configs['train_input_config'] model_fn = functools.partial(model_builder.build, model_config=model_config, is_training=True) def get_next(config): return dataset_util.make_initializable_iterator( dataset_builder.build(config)).get_next() create_input_dict_fn = functools.partial(get_next, input_config) 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)
pipeline_config = os.path.join(CWD_PATH, MODEL_NAME, 'pipeline.config') # Number of classes the object detector can identify NUM_CLASSES = 90 ## Load the label map. # Label maps map indices to category names, so that when the convolution # network predicts `5`, we know that this corresponds to `airplane`. # Here we use internal utility functions, but anything that returns a # dictionary mapping integers to appropriate string labels would be fine label_map = label_map_util.load_labelmap(PATH_TO_LABELS) categories = label_map_util.convert_label_map_to_categories( label_map, max_num_classes=NUM_CLASSES, use_display_name=True) category_index = label_map_util.create_category_index(categories) configs = config_util.get_configs_from_pipeline_file(pipeline_config) detection_model = model_builder.build(model_config=model_config, is_training=False) ckpt = tf.compat.v2.train.Checkpoint(model=detection_model) ckpt.restore(os.path.join(model_dir, 'ckpt-0')).expect_partial() def get_model_detection_function(model): """Get a tf.function for detection.""" @tf.function def detect_fn(image): """Detect objects in image.""" image, shapes = model.preprocess(image) prediction_dict = model.predict(image, shapes) detections = model.postprocess(prediction_dict, shapes)
gpu_options=tf.compat.v1.GPUOptions(per_process_gpu_memory_fraction=0.8) # device_count = {'GPU': 1} ) config.gpu_options.allow_growth = True session = tf.compat.v1.Session(config=config) tf.compat.v1.keras.backend.set_session(session) from utils import config_util from builders import model_builder CONFIG_PATH = 'new_graph/pipeline.config' CHECKPOINT_PATH = 'new_graph/checkpoint/ckpt-0' ANNOTATION_PATH = 'label_map.pbtxt' # Load pipeline config and build a detection model configs = config_util.get_configs_from_pipeline_file(CONFIG_PATH) detection_model = model_builder.build(model_config=configs['model'], is_training=False) # Restore checkpoint ckpt = tf.compat.v2.train.Checkpoint(model=detection_model) ckpt.restore(CHECKPOINT_PATH).expect_partial() @tf.function def detect_fn(image): image, shapes = detection_model.preprocess(image) prediction_dict = detection_model.predict(image, shapes) detections = detection_model.postprocess(prediction_dict, shapes) return detections
def create_estimator(run_config, hparams, pipeline_config_path, train_steps=None, eval_steps=None, train_batch_size=None, model_fn_creator=model.create_model_fn, use_tpu=False, num_shards=1, params=None, **kwargs): """Creates an `Estimator` object. Args: run_config: A `RunConfig`. hparams: A `HParams`. pipeline_config_path: A path to a pipeline config file. train_steps: Number of training steps. If None, the number of training steps is set from the `TrainConfig` proto. eval_steps: Number of evaluation steps per evaluation cycle. If None, the number of evaluation steps is set from the `EvalConfig` proto. train_batch_size: Training batch size. If none, use batch size from `TrainConfig` proto. model_fn_creator: A function that creates a `model_fn` for `Estimator`. Follows the signature: * Args: * `detection_model_fn`: Function that returns `DetectionModel` instance. * `configs`: Dictionary of pipeline config objects. * `hparams`: `HParams` object. * Returns: `model_fn` for `Estimator`. use_tpu: Boolean, whether training and evaluation should run on TPU. num_shards: Number of shards (TPU cores). params: Parameter dictionary passed from the estimator. **kwargs: Additional keyword arguments for configuration override. Returns: Estimator: A estimator object used for training and evaluation train_input_fn: Input function for the training loop eval_input_fn: Input function for the evaluation run train_steps: Number of training steps either from arg `train_steps` or `TrainConfig` proto eval_steps: Number of evaluation steps either from arg `eval_steps` or `EvalConfig` proto """ configs = config_util.get_configs_from_pipeline_file(pipeline_config_path) configs = config_util.merge_external_params_with_configs( configs, hparams, train_steps=train_steps, eval_steps=eval_steps, batch_size=train_batch_size, **kwargs) model_config = configs['model'] train_config = configs['train_config'] train_input_config = configs['train_input_config'] eval_config = configs['eval_config'] eval_input_config = configs['eval_input_config'] if params is None: params = {} if train_steps is None: train_steps = train_config.num_steps if train_config.num_steps else None if eval_steps is None: eval_steps = eval_config.num_examples if eval_config.num_examples else None detection_model_fn = functools.partial( model_builder.build, model_config=model_config) # Create the input functions for TRAIN/EVAL. train_input_fn = inputs.create_train_input_fn( train_config=train_config, train_input_config=train_input_config, model_config=model_config) eval_input_fn = inputs.create_eval_input_fn( eval_config=eval_config, eval_input_config=eval_input_config, model_config=model_config) estimator = tpu_estimator.TPUEstimator( model_fn=model_fn_creator(detection_model_fn, configs, hparams, use_tpu), train_batch_size=train_config.batch_size, # For each core, only batch size 1 is supported for eval. eval_batch_size=num_shards * 1 if use_tpu else 1, use_tpu=use_tpu, config=run_config, params=params) return estimator, train_input_fn, eval_input_fn, train_steps, eval_steps
def _assertOptimizerWithNewLearningRate(self, optimizer_name): """Asserts successful updating of all learning rate schemes.""" original_learning_rate = 0.7 learning_rate_scaling = 0.1 hparams = tf.contrib.training.HParams(learning_rate=0.15) pipeline_config_path = os.path.join(self.get_temp_dir(), "pipeline.config") # Constant learning rate. pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() optimizer = getattr(pipeline_config.train_config.optimizer, optimizer_name) _update_optimizer_with_constant_learning_rate(optimizer, original_learning_rate) _write_config(pipeline_config, pipeline_config_path) configs = config_util.get_configs_from_pipeline_file( pipeline_config_path) configs = config_util.merge_external_params_with_configs( configs, hparams) optimizer = getattr(configs["train_config"].optimizer, optimizer_name) constant_lr = optimizer.learning_rate.constant_learning_rate self.assertAlmostEqual(hparams.learning_rate, constant_lr.learning_rate) # Exponential decay learning rate. pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() optimizer = getattr(pipeline_config.train_config.optimizer, optimizer_name) _update_optimizer_with_exponential_decay_learning_rate( optimizer, original_learning_rate) _write_config(pipeline_config, pipeline_config_path) configs = config_util.get_configs_from_pipeline_file( pipeline_config_path) configs = config_util.merge_external_params_with_configs( configs, hparams) optimizer = getattr(configs["train_config"].optimizer, optimizer_name) exponential_lr = optimizer.learning_rate.exponential_decay_learning_rate self.assertAlmostEqual(hparams.learning_rate, exponential_lr.initial_learning_rate) # Manual step learning rate. pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() optimizer = getattr(pipeline_config.train_config.optimizer, optimizer_name) _update_optimizer_with_manual_step_learning_rate( optimizer, original_learning_rate, learning_rate_scaling) _write_config(pipeline_config, pipeline_config_path) configs = config_util.get_configs_from_pipeline_file( pipeline_config_path) configs = config_util.merge_external_params_with_configs( configs, hparams) optimizer = getattr(configs["train_config"].optimizer, optimizer_name) manual_lr = optimizer.learning_rate.manual_step_learning_rate self.assertAlmostEqual(hparams.learning_rate, manual_lr.initial_learning_rate) for i, schedule in enumerate(manual_lr.schedule): self.assertAlmostEqual( hparams.learning_rate * learning_rate_scaling**i, schedule.learning_rate)