def argument_parser() -> argparse.Namespace:
parser = argparse.ArgumentParser()
exclusive_args = parser.add_mutually_exclusive_group()
exclusive_args.add_argument(
"--model",
type=str,
default=None,
help="Saved model location to use for inference")
exclusive_args.add_argument("--architecture",
type=str,
choices=model_architectures.keys())
parser.add_argument("--log-path",
type=str,
default="./log.json",
help="Path to log file")
parser.add_argument("--tf-trt",
action="store_true",
default=False,
help="Use TF-TRT for inference")
parser.add_argument("--amp",
action="store_true",
default=False,
help="Use AMP for inference")
parser.add_argument("--data-dir",
type=str,
required=False,
default=None,
help="Localization of validation data")
parser.add_argument("--batch-size",
type=int,
default=1,
help="Batch size for inference")
return parser.parse_args()
import tensorflow as tf
import horovod.tensorflow as hvd
import dllogger
from utils import hvd_utils
from runtime import Runner
from model.resnet import model_architectures
from utils.cmdline_helper import parse_cmdline
if __name__ == "__main__":
tf.logging.set_verbosity(tf.logging.ERROR)
FLAGS = parse_cmdline(model_architectures.keys())
hvd.init()
if hvd.rank() == 0:
log_path = os.path.join(FLAGS.results_dir, FLAGS.log_filename)
os.makedirs(FLAGS.results_dir, exist_ok=True)
dllogger.init(backends=[
dllogger.JSONStreamBackend(verbosity=dllogger.Verbosity.VERBOSE,
filename=log_path),
dllogger.StdOutBackend(verbosity=dllogger.Verbosity.VERBOSE)
])
else:
dllogger.init(backends=[])
dllogger.log(data=vars(FLAGS), step='PARAMETER')