parser.add_argument( '--config', help="A list of KEY=VALUE to overwrite those defined in config.py", nargs='+') if get_tf_version_tuple() < (1, 6): # https://github.com/tensorflow/tensorflow/issues/14657 logger.warn( "TF<1.6 has a bug which may lead to crash in FasterRCNN if you're unlucky." ) args = parser.parse_args() if args.config: cfg.update_args(args.config) register_coco(cfg.DATA.BASEDIR) # add COCO datasets to the registry register_balloon( cfg.DATA.BASEDIR) # add the demo balloon datasets to the registry # Setup logging ... is_horovod = cfg.TRAINER == 'horovod' if is_horovod: hvd.init() if not is_horovod or hvd.rank() == 0: logger.set_logger_dir(args.logdir, 'd') logger.info("Environment Information:\n" + collect_env_info()) finalize_configs(is_training=True) # Create model MODEL = ResNetFPNModel() if cfg.MODE_FPN else ResNetC4Model() # Compute the training schedule from the number of GPUs ...
parser.add_argument( '--benchmark', action='store_true', help="Benchmark the speed of the model + postprocessing") parser.add_argument( '--config', help="A list of KEY=VALUE to overwrite those defined in config.py", nargs='+') parser.add_argument('--compact', help='Save a model to .pb') parser.add_argument('--serving', help='Save a model to serving file') args = parser.parse_args() if args.config: cfg.update_args(args.config) register_coco(cfg.DATA.BASEDIR) # add COCO datasets to the registry register_balloon(cfg.DATA.BASEDIR) MODEL = ResNetFPNModel() if cfg.MODE_FPN else ResNetC4Model() if not tf.test.is_gpu_available(): from tensorflow.python.framework import test_util assert get_tf_version_tuple() >= (1, 7) and test_util.IsMklEnabled(), \ "Inference requires either GPU support or MKL support!" assert args.load finalize_configs(is_training=False) if args.predict or args.visualize: cfg.TEST.RESULT_SCORE_THRESH = cfg.TEST.RESULT_SCORE_THRESH_VIS if args.visualize: do_visualize(MODEL, args.load)