def evaluation_data_fn():
if self.run_hparams.data_dir is not None:
return data_utils.get_tfrecords_input_fn(
filenames=filenames,
batch_size=batch_size,
height=self.run_hparams.height,
width=self.run_hparams.width,
training=False,
distort_color=self.run_hparams.distort_colors,
num_threads=self.run_hparams.num_preprocessing_threads,
deterministic=False
if self.run_hparams.seed is None else True)
else:
LOGGER.log("Using Synthetic Data ...\n")
return data_utils.get_synth_input_fn(
batch_size=batch_size,
height=self.run_hparams.height,
width=self.run_hparams.width,
num_channels=self.run_hparams.n_channels,
data_format=self.run_hparams.input_format,
num_classes=self.run_hparams.n_classes,
dtype=self.run_hparams.dtype,
)
def training_data_fn():
if self.run_hparams.data_dir is not None:
return data_utils.get_tfrecords_input_fn(
data_dir=self.run_hparams.data_dir,
num_epochs=num_iter,
batch_size=batch_size,
height=self.run_hparams.height,
width=self.run_hparams.width,
training=True,
# distort_color=self.run_hparams.distort_colors,
# num_threads=self.run_hparams.num_preprocessing_threads,
datasets_num_private_threads=None
# deterministic=False if self.run_hparams.seed is None else True
)
else:
if hvd.rank() == 0:
LOGGER.log("Using Synthetic Data ...")
return data_utils.get_synth_input_fn(
batch_size=batch_size,
height=self.run_hparams.height,
width=self.run_hparams.width,
num_channels=self.run_hparams.n_channels,
data_format=self.run_hparams.input_format,
num_classes=self.run_hparams.n_classes,
dtype=self.run_hparams.dtype,
)
def training_data_fn():
if self.run_hparams.use_dali and self.run_hparams.data_idx_dir is not None:
if hvd.rank() == 0:
print("Using DALI input... ")
return data_utils.get_dali_input_fn(
filenames=filenames,
idx_filenames=idx_filenames,
batch_size=batch_size,
height=self.run_hparams.height,
width=self.run_hparams.width,
training=True,
distort_color=self.run_hparams.distort_colors,
num_threads=self.run_hparams.num_preprocessing_threads,
deterministic=False
if self.run_hparams.seed is None else True)
elif self.run_hparams.data_dir is not None:
return data_utils.get_tfrecords_input_fn(
filenames=filenames,
batch_size=batch_size,
height=self.run_hparams.height,
width=self.run_hparams.width,
training=True,
distort_color=self.run_hparams.distort_colors,
num_threads=self.run_hparams.num_preprocessing_threads,
deterministic=False
if self.run_hparams.seed is None else True)
else:
if hvd.rank() == 0:
print("Using Synthetic Data ...")
return data_utils.get_synth_input_fn(
batch_size=batch_size,
height=self.run_hparams.height,
width=self.run_hparams.width,
num_channels=self.run_hparams.n_channels,
data_format=self.run_hparams.input_format,
num_classes=self.run_hparams.n_classes,
dtype=self.run_hparams.dtype,
)
def evaluation_data_fn():
if self.run_hparams.data_dir is not None:
return data_utils.get_tfrecords_input_fn(
data_dir=self.run_hparams.data_dir,
num_epochs=num_iter,
batch_size=batch_size,
height=self.run_hparams.height,
width=self.run_hparams.width,
training=True,
datasets_num_private_threads=None)
else:
LOGGER.log("Using Synthetic Data ...\n")
return data_utils.get_synth_input_fn(
batch_size=batch_size,
height=self.run_hparams.height,
width=self.run_hparams.width,
num_channels=self.run_hparams.n_channels,
data_format=self.run_hparams.input_format,
num_classes=self.run_hparams.n_classes,
dtype=self.run_hparams.dtype,
)
def main(args: argparse.Namespace):
hvd.init()
dllogger.init(backends=[
dllogger.JSONStreamBackend(verbosity=dllogger.Verbosity.VERBOSE,
filename=args.log_path),
dllogger.StdOutBackend(verbosity=dllogger.Verbosity.VERBOSE)
])
dllogger.log(data=vars(args), step='PARAMETER')
if args.model is None:
saved_model_to_load = tempfile.mkdtemp(prefix="tftrt-savedmodel")
r = runner.Runner(n_classes=1001,
architecture=args.architecture,
use_tf_amp=args.amp,
model_dir=saved_model_to_load)
r.train("batch", 1, 1, args.batch_size, is_benchmark=True)
r.evaluate("batch",
1,
args.batch_size,
export_dir=saved_model_to_load,
is_benchmark=True)
saved_model_to_load = r.exported_path.decode("utf-8")
else:
saved_model_to_load = args.model
output_tensor_name = "y_preds_ref:0" if not args.tf_trt else "ArgMax:0"
batch_size = args.batch_size
if args.tf_trt:
converter = trt.TrtGraphConverter(
input_saved_model_dir=str(saved_model_to_load),
precision_mode="FP16" if args.amp else "FP32")
converter.convert()
converter.save(OUTPUT_SAVED_MODEL_PATH)
saved_model_to_load = OUTPUT_SAVED_MODEL_PATH
elif args.amp:
os.environ["TF_ENABLE_AUTO_MIXED_PRECISION_GRAPH_REWRITE"] = "1"
if args.data_dir is not None:
filenames, _, num_steps, _, _ = runner_utils.parse_tfrecords_dataset(
data_dir=str(args.data_dir),
mode="validation",
iter_unit="epoch",
num_iter=1,
global_batch_size=batch_size,
)
dataset = data_utils.get_tfrecords_input_fn(filenames=filenames,
batch_size=batch_size,
height=224,
width=224,
training=False,
distort_color=False,
num_threads=1,
deterministic=True)
iterator = dataset.make_initializable_iterator()
next_item = iterator.get_next()
else:
num_steps = 60000 / batch_size
with tf.Session() as sess:
if args.data_dir is not None:
sess.run(iterator.initializer)
tf.saved_model.loader.load(sess,
[tf.saved_model.tag_constants.SERVING],
str(saved_model_to_load))
try:
start_time = time.time()
last_time = start_time
image_processed = 0
image_correct = 0
for samples_processed in range(int(num_steps)):
if args.data_dir is not None:
next_batch_image, next_batch_target = sess.run(next_item)
else:
if samples_processed == 0:
next_batch_image = np.random.normal(size=(batch_size,
224, 224, 3))
next_batch_target = np.random.randint(
0, 1000, size=(batch_size, ))
output = sess.run(
[output_tensor_name],
feed_dict={"input_tensor:0": next_batch_image})
image_processed += args.batch_size
image_correct += np.sum(output == next_batch_target)
if samples_processed % LOG_FREQUENCY == 0 and samples_processed != 0:
current_time = time.time()
current_throughput = LOG_FREQUENCY * batch_size / (
current_time - last_time)
dllogger.log(step=(0, samples_processed),
data={"throughput": current_throughput})
last_time = current_time
except tf.errors.OutOfRangeError:
pass
finally:
dllogger.log(step=tuple(),
data={
"throughput":
image_processed / (last_time - start_time),
"accuracy":
image_correct / image_processed
})
