def _preprocessing(self, raw_images, image_size, mode=None):
"""Preprocess images before feeding to the network."""
if not mode:
return raw_images, None
image_size = model_utils.parse_image_size(image_size)
if mode != 'infer':
# We only support inference for now.
raise ValueError('preprocessing must be infer or empty')
def map_fn(image):
input_processor = dataloader.DetectionInputProcessor(
image, image_size)
input_processor.normalize_image()
input_processor.set_scale_factors_to_output_size()
image = input_processor.resize_and_crop_image()
image_scale = input_processor.image_scale_to_original
return image, image_scale
if raw_images.shape.as_list()[0]: # fixed batch size.
batch_size = raw_images.shape.as_list()[0]
outputs = [map_fn(raw_images[i]) for i in range(batch_size)]
return [tf.stack(y) for y in zip(*outputs)]
# otherwise treat it as dynamic batch size.
return tf.vectorized_map(map_fn, raw_images)
def __init__(self, min_level, max_level, num_scales, aspect_ratios,
anchor_scale, image_size):
"""Constructs multiscale RetinaNet anchors.
Args:
min_level: integer number of minimum level of the output feature pyramid.
max_level: integer number of maximum level of the output feature pyramid.
num_scales: integer number representing intermediate scales added
on each level. For instances, num_scales=2 adds two additional
anchor scales [2^0, 2^0.5] on each level.
aspect_ratios: list of tuples representing the aspect ratio anchors added
on each level. For instances, aspect_ratios =
[(1, 1), (1.4, 0.7), (0.7, 1.4)] adds three anchors on each level.
anchor_scale: float number representing the scale of size of the base
anchor to the feature stride 2^level.
image_size: integer number or tuple of integer number of input image size.
"""
self.min_level = min_level
self.max_level = max_level
self.num_scales = num_scales
self.aspect_ratios = aspect_ratios
self.anchor_scale = anchor_scale
self.image_size = model_utils.parse_image_size(image_size)
self.feat_sizes = model_utils.get_feat_sizes(image_size, max_level)
self.config = self._generate_configs()
self.boxes = self._generate_boxes()
def build(self, params_override=None):
"""Build model and restore checkpoints."""
params = copy.deepcopy(self.params)
if params_override:
params.update(params_override)
config = hparams_config.get_efficientdet_config(self.model_name)
config.override(params)
self.model = efficientdet_keras.EfficientDetModel(config=config)
image_size = model_utils.parse_image_size(params['image_size'])
self.model.build((self.batch_size, *image_size, 3))
util_keras.restore_ckpt(self.model, self.ckpt_path, 0,
params['moving_average_decay'])
def set_model(self, model: tf.keras.Model):
self.train_model = model
with tf.device('/cpu:0'):
self.model = efficientdet_keras.EfficientDetModel(
config=model.config)
height, width = model_utils.parse_image_size(model.config.image_size)
self.model.build((1, height, width, 3))
self.file_writer = tf.summary.create_file_writer(self.output_dir)
self.min_score_thresh = self.model.config.nms_configs[
'score_thresh'] or 0.4
self.max_boxes_to_draw = (
self.model.config.nms_configs['max_output_size'] or 100)
def set_training_random_scale_factors(self,
scale_min,
scale_max,
target_size=None):
"""Set the parameters for multiscale training.
Notably, if train and eval use different sizes, then target_size should be
set as eval size to avoid the discrency between train and eval.
Args:
scale_min: minimal scale factor.
scale_max: maximum scale factor.
target_size: targeted size, usually same as eval. If None, use train size.
"""
if not target_size:
target_size = self._output_size
target_size = model_utils.parse_image_size(target_size)
logging.info('target_size = %s, output_size = %s', target_size,
self._output_size)
# Select a random scale factor.
random_scale_factor = tf.random.uniform([], scale_min, scale_max)
scaled_y = tf.cast(random_scale_factor * target_size[0], tf.int32)
scaled_x = tf.cast(random_scale_factor * target_size[1], tf.int32)
# Recompute the accurate scale_factor using rounded scaled image size.
height = tf.cast(tf.shape(self._image)[0], tf.float32)
width = tf.cast(tf.shape(self._image)[1], tf.float32)
image_scale_y = tf.cast(scaled_y, tf.float32) / height
image_scale_x = tf.cast(scaled_x, tf.float32) / width
image_scale = tf.minimum(image_scale_x, image_scale_y)
# Select non-zero random offset (x, y) if scaled image is larger than
# self._output_size.
scaled_height = tf.cast(height * image_scale, tf.int32)
scaled_width = tf.cast(width * image_scale, tf.int32)
offset_y = tf.cast(scaled_height - self._output_size[0], tf.float32)
offset_x = tf.cast(scaled_width - self._output_size[1], tf.float32)
offset_y = tf.maximum(0.0, offset_y) * tf.random.uniform([], 0, 1)
offset_x = tf.maximum(0.0, offset_x) * tf.random.uniform([], 0, 1)
offset_y = tf.cast(offset_y, tf.int32)
offset_x = tf.cast(offset_x, tf.int32)
self._image_scale = image_scale
self._scaled_height = scaled_height
self._scaled_width = scaled_width
self._crop_offset_x = offset_x
self._crop_offset_y = offset_y
def main(_):
# get e2e training time
begin = time.time()
logging.info("Training started at: {}".format(time.asctime()))
hvd.init()
# Parse and override hparams
config = hparams_config.get_detection_config(FLAGS.model_name)
config.override(FLAGS.hparams)
if FLAGS.num_epochs: # NOTE: remove this flag after updating all docs.
config.num_epochs = FLAGS.num_epochs
if FLAGS.lr:
config.learning_rate = FLAGS.lr
if FLAGS.warmup_value:
config.lr_warmup_init = FLAGS.warmup_value
if FLAGS.warmup_epochs:
config.lr_warmup_epoch = FLAGS.warmup_epochs
config.backbone_init = FLAGS.backbone_init
config.mixed_precision = FLAGS.amp
config.image_size = model_utils.parse_image_size(config.image_size)
# get eval config
eval_config = hparams_config.get_detection_config(FLAGS.model_name)
eval_config.override(FLAGS.hparams)
eval_config.val_json_file = FLAGS.val_json_file
eval_config.val_file_pattern = FLAGS.val_file_pattern
eval_config.nms_configs.max_nms_inputs = anchors.MAX_DETECTION_POINTS
eval_config.drop_remainder = False # eval all examples w/o drop.
eval_config.image_size = model_utils.parse_image_size(
eval_config['image_size'])
# setup
setup.set_flags(FLAGS, config, training=True)
if FLAGS.debug:
tf.config.experimental_run_functions_eagerly(True)
tf.debugging.set_log_device_placement(True)
tf.random.set_seed(111111)
logging.set_verbosity(logging.DEBUG)
# Check data path
if FLAGS.training_file_pattern is None or FLAGS.val_file_pattern is None or FLAGS.val_json_file is None:
raise RuntimeError(
'You must specify --training_file_pattern, --val_file_pattern and --val_json_file for training.'
)
steps_per_epoch = (FLAGS.num_examples_per_epoch +
(FLAGS.batch_size * get_world_size()) -
1) // (FLAGS.batch_size * get_world_size())
if FLAGS.benchmark == True:
# For ci perf training runs, run for a fixed number of iterations per epoch
steps_per_epoch = FLAGS.benchmark_steps
params = dict(config.as_dict(),
model_name=FLAGS.model_name,
model_dir=FLAGS.model_dir,
steps_per_epoch=steps_per_epoch,
checkpoint_period=FLAGS.checkpoint_period,
batch_size=FLAGS.batch_size,
num_shards=get_world_size(),
val_json_file=FLAGS.val_json_file,
testdev_dir=FLAGS.testdev_dir,
mode='train')
logging.info('Training params: {}'.format(params))
# make output dir if it does not exist
tf.io.gfile.makedirs(FLAGS.model_dir)
# dllogger setup
backends = []
if is_main_process():
log_path = os.path.join(FLAGS.model_dir, FLAGS.log_filename)
backends += [
JSONStreamBackend(verbosity=Verbosity.VERBOSE, filename=log_path),
StdOutBackend(verbosity=Verbosity.DEFAULT)
]
DLLogger.init(backends=backends)
def get_dataset(is_training, params):
file_pattern = (FLAGS.training_file_pattern
if is_training else FLAGS.val_file_pattern)
if not file_pattern:
raise ValueError('No matching files.')
return dataloader.InputReader(
file_pattern,
is_training=is_training,
use_fake_data=FLAGS.use_fake_data,
max_instances_per_image=config.max_instances_per_image,
enable_map_parallelization=FLAGS.enable_map_parallelization)(
params)
num_samples = (FLAGS.eval_samples + get_world_size() -
1) // get_world_size()
num_samples = (num_samples + FLAGS.eval_batch_size -
1) // FLAGS.eval_batch_size
eval_config.num_samples = num_samples
def get_eval_dataset(eval_config):
dataset = dataloader.InputReader(
FLAGS.val_file_pattern,
is_training=False,
max_instances_per_image=eval_config.max_instances_per_image)(
eval_config, batch_size=FLAGS.eval_batch_size)
dataset = dataset.shard(get_world_size(), get_rank())
dataset = dataset.take(num_samples)
return dataset
eval_dataset = get_eval_dataset(eval_config)
# pick focal loss implementation
focal_loss = train_lib.StableFocalLoss(
params['alpha'],
params['gamma'],
label_smoothing=params['label_smoothing'],
reduction=tf.keras.losses.Reduction.NONE)
model = train_lib.EfficientDetNetTrain(params['model_name'], config)
model.build((None, *config.image_size, 3))
model.compile(
optimizer=optimizer_builder.get_optimizer(params),
loss={
'box_loss':
train_lib.BoxLoss(params['delta'],
reduction=tf.keras.losses.Reduction.NONE),
'box_iou_loss':
train_lib.BoxIouLoss(params['iou_loss_type'],
params['min_level'],
params['max_level'],
params['num_scales'],
params['aspect_ratios'],
params['anchor_scale'],
params['image_size'],
reduction=tf.keras.losses.Reduction.NONE),
'class_loss':
focal_loss,
'seg_loss':
tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True, reduction=tf.keras.losses.Reduction.NONE)
})
train_from_epoch = util_keras.restore_ckpt(model,
params['model_dir'],
config.moving_average_decay,
steps_per_epoch=steps_per_epoch)
print("training_mode: {}".format(FLAGS.training_mode))
callbacks = callback_builder.get_callbacks(params, FLAGS.training_mode,
eval_config, eval_dataset,
DLLogger, FLAGS.time_history,
FLAGS.log_steps, FLAGS.lr_tb,
FLAGS.benchmark)
history = model.fit(
get_dataset(True, params=params),
epochs=params['num_epochs'],
steps_per_epoch=steps_per_epoch,
initial_epoch=train_from_epoch,
callbacks=callbacks,
verbose=1 if is_main_process() else 0,
validation_data=get_dataset(False, params=params)
if FLAGS.validate else None,
validation_steps=(FLAGS.eval_samples //
FLAGS.eval_batch_size) if FLAGS.validate else None)
if is_main_process():
model.save_weights(os.path.join(FLAGS.model_dir, 'ckpt-final'))
# log final stats
stats = {}
for callback in callbacks:
if isinstance(callback, callback_builder.TimeHistory):
if callback.epoch_runtime_log:
stats[
'avg_fps_training'] = callback.average_examples_per_second
stats[
'avg_fps_training_per_GPU'] = callback.average_examples_per_second / get_world_size(
)
stats[
'avg_latency_training'] = callback.average_time_per_iteration
if history and history.history:
train_hist = history.history
#Gets final loss from training.
stats['training_loss'] = float(
hvd.allreduce(tf.constant(train_hist['loss'][-1],
dtype=tf.float32),
average=True))
if os.path.exists(os.path.join(FLAGS.model_dir, 'ema_weights')):
ckpt_epoch = "%02d" % sorted(set([
int(f.rsplit('.')[0].rsplit('-')[1])
for f in os.listdir(os.path.join(FLAGS.model_dir, 'ema_weights'))
if 'emackpt' in f
]),
reverse=True)[0]
ckpt = os.path.join(FLAGS.model_dir, 'ema_weights',
'emackpt-' + str(ckpt_epoch))
util_keras.restore_ckpt(model,
ckpt,
eval_config.moving_average_decay,
steps_per_epoch=0,
skip_mismatch=False,
expect_partial=True)
if is_main_process():
model.save(os.path.join(FLAGS.model_dir, 'emackpt-final'))
else:
ckpt_epoch = 'final'
ckpt = os.path.join(FLAGS.model_dir, 'ckpt-' + ckpt_epoch)
if is_main_process():
model.save(os.path.join(FLAGS.model_dir, 'ckpt-' + ckpt_epoch))
# Start evaluation of final ema checkpoint
logging.set_verbosity(logging.WARNING)
@tf.function
def model_fn(images, labels):
cls_outputs, box_outputs = model(images, training=False)
detections = postprocess.generate_detections(eval_config, cls_outputs,
box_outputs,
labels['image_scales'],
labels['source_ids'])
tf.numpy_function(evaluator.update_state, [
labels['groundtruth_data'],
postprocess.transform_detections(detections)
], [])
if FLAGS.benchmark == False and FLAGS.training_mode == 'train':
# Evaluator for AP calculation.
label_map = label_util.get_label_map(eval_config.label_map)
evaluator = coco_metric.EvaluationMetric(
filename=eval_config.val_json_file, label_map=label_map)
evaluator.reset_states()
# evaluate all images.
pbar = tf.keras.utils.Progbar(num_samples)
for i, (images, labels) in enumerate(eval_dataset):
model_fn(images, labels)
if is_main_process():
pbar.update(i)
# gather detections from all ranks
evaluator.gather()
if is_main_process():
# compute the final eval results.
metrics = evaluator.result()
metric_dict = {}
for i, name in enumerate(evaluator.metric_names):
metric_dict[name] = metrics[i]
if label_map:
for i, cid in enumerate(sorted(label_map.keys())):
name = 'AP_/%s' % label_map[cid]
metric_dict[name] = metrics[i +
len(evaluator.metric_names)]
# csv format
csv_metrics = ['AP', 'AP50', 'AP75', 'APs', 'APm', 'APl']
csv_format = ",".join(
[str(ckpt_epoch)] +
[str(round(metric_dict[key] * 100, 2)) for key in csv_metrics])
print(FLAGS.model_name, metric_dict, "csv format:", csv_format)
MPI.COMM_WORLD.Barrier()
if is_main_process():
stats['e2e_training_time'] = time.time() - begin
DLLogger.log(step=(), data=stats)
def main(_):
hvd.init()
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
if gpus:
tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')
if FLAGS.amp:
policy = tf.keras.mixed_precision.experimental.Policy('mixed_float16')
tf.keras.mixed_precision.experimental.set_policy(policy)
else:
os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '0'
config = hparams_config.get_efficientdet_config(FLAGS.model_name)
config.override(FLAGS.hparams)
config.val_json_file = FLAGS.val_json_file
config.nms_configs.max_nms_inputs = anchors.MAX_DETECTION_POINTS
config.drop_remainder = False # eval all examples w/o drop.
config.image_size = model_utils.parse_image_size(config['image_size'])
@tf.function
def model_fn(images, labels):
cls_outputs, box_outputs = model(images, training=False)
detections = postprocess.generate_detections(config, cls_outputs, box_outputs,
labels['image_scales'],
labels['source_ids'])
tf.numpy_function(evaluator.update_state,
[labels['groundtruth_data'],
postprocess.transform_detections(detections)], [])
# Network
model = efficientdet_keras.EfficientDetNet(config=config)
model.build((None, *config.image_size, 3))
# dataset
batch_size = FLAGS.batch_size # local batch size.
ds = dataloader.InputReader(
FLAGS.val_file_pattern,
is_training=False,
max_instances_per_image=config.max_instances_per_image,
enable_map_parallelization=FLAGS.enable_map_parallelization)(
config, batch_size=batch_size)
ds = ds.shard(get_world_size(), get_rank())
# Evaluator for AP calculation.
label_map = label_util.get_label_map(config.label_map)
evaluator = coco_metric.EvaluationMetric(
filename=config.val_json_file, label_map=label_map)
util_keras.restore_ckpt(model, FLAGS.ckpt_path, config.moving_average_decay,
steps_per_epoch=0, skip_mismatch=False, expect_partial=True)
if FLAGS.eval_samples:
num_samples = (FLAGS.eval_samples + get_world_size() - 1) // get_world_size()
num_samples = (num_samples + batch_size - 1) // batch_size
ds = ds.take(num_samples)
evaluator.reset_states()
# evaluate all images.
pbar = tf.keras.utils.Progbar(num_samples)
for i, (images, labels) in enumerate(ds):
model_fn(images, labels)
if is_main_process():
pbar.update(i)
# gather detections from all ranks
evaluator.gather()
if is_main_process():
# compute the final eval results.
metrics = evaluator.result()
metric_dict = {}
for i, name in enumerate(evaluator.metric_names):
metric_dict[name] = metrics[i]
if label_map:
for i, cid in enumerate(sorted(label_map.keys())):
name = 'AP_/%s' % label_map[cid]
metric_dict[name] = metrics[i + len(evaluator.metric_names)]
# csv format
csv_metrics = ['AP','AP50','AP75','APs','APm','APl']
csv_format = ",".join([str(round(metric_dict[key] * 100, 2)) for key in csv_metrics])
print(FLAGS.model_name, metric_dict, "csv format:", csv_format)
MPI.COMM_WORLD.Barrier()
def main(_):
model_config = hparams_config.get_detection_config(FLAGS.model_name)
model_config.override(FLAGS.hparams) # Add custom overrides
model_config.is_training_bn = False
model_config.image_size = model_utils.parse_image_size(model_config.image_size)
# A hack to make flag consistent with nms configs.
if FLAGS.min_score_thresh:
model_config.nms_configs.score_thresh = FLAGS.min_score_thresh
if FLAGS.nms_method:
model_config.nms_configs.method = FLAGS.nms_method
if FLAGS.max_boxes_to_draw:
model_config.nms_configs.max_output_size = FLAGS.max_boxes_to_draw
model_config.mixed_precision = FLAGS.amp
setup.set_flags(FLAGS, model_config, training=False)
model_params = model_config.as_dict()
ckpt_path_or_file = FLAGS.ckpt_path
if tf.io.gfile.isdir(ckpt_path_or_file):
ckpt_path_or_file = tf.train.latest_checkpoint(ckpt_path_or_file)
driver = inference.ServingDriver(FLAGS.model_name, ckpt_path_or_file,
FLAGS.batch_size or None,
FLAGS.min_score_thresh,
FLAGS.max_boxes_to_draw, model_params)
# dllogger setup
backends = []
backends+=[
JSONStreamBackend(verbosity=Verbosity.VERBOSE, filename=FLAGS.dllogger_path),
StdOutBackend(verbosity=Verbosity.DEFAULT)]
DLLogger.init(backends=backends)
if FLAGS.mode == 'export':
if tf.io.gfile.exists(FLAGS.saved_model_dir):
tf.io.gfile.rmtree(FLAGS.saved_model_dir)
driver.export(FLAGS.saved_model_dir, FLAGS.tflite_path, FLAGS.tensorrt)
elif FLAGS.mode == 'benchmark':
if FLAGS.saved_model_dir:
driver.load(FLAGS.saved_model_dir)
batch_size = FLAGS.batch_size or 1
if FLAGS.input_image:
image_file = tf.io.read_file(FLAGS.input_image)
image_arrays = tf.image.decode_image(image_file)
image_arrays.set_shape((None, None, 3))
image_arrays = tf.expand_dims(image_arrays, 0)
if batch_size > 1:
image_arrays = tf.tile(image_arrays, [batch_size, 1, 1, 1])
else:
# use synthetic data if no image is provided.
image_arrays = tf.ones((batch_size, *model_config.image_size, 3),
dtype=tf.uint8)
driver.benchmark(image_arrays, FLAGS.bm_runs, FLAGS.trace_filename)
elif FLAGS.mode == 'dry':
# transfer to tf2 format ckpt
driver.build()
if FLAGS.export_ckpt:
driver.model.save_weights(FLAGS.export_ckpt)
elif FLAGS.mode == 'video':
import cv2 # pylint: disable=g-import-not-at-top
if tf.saved_model.contains_saved_model(FLAGS.saved_model_dir):
driver.load(FLAGS.saved_model_dir)
cap = cv2.VideoCapture(FLAGS.input_video)
if not cap.isOpened():
print('Error opening input video: {}'.format(FLAGS.input_video))
out_ptr = None
if FLAGS.output_video:
frame_width, frame_height = int(cap.get(3)), int(cap.get(4))
out_ptr = cv2.VideoWriter(FLAGS.output_video,
cv2.VideoWriter_fourcc('m', 'p', '4', 'v'), 25,
(frame_width, frame_height))
while cap.isOpened():
# Capture frame-by-frame
ret, frame = cap.read()
if not ret:
break
raw_frames = np.array([frame])
detections_bs = driver.serve(raw_frames)
boxes, scores, classes, _ = tf.nest.map_structure(np.array, detections_bs)
new_frame = driver.visualize(
raw_frames[0],
boxes[0],
scores[0],
classes[0],
min_score_thresh=model_config.nms_configs.score_thresh,
max_boxes_to_draw=model_config.nms_configs.max_output_size)
if out_ptr:
# write frame into output file.
out_ptr.write(new_frame)
else:
# show the frame online, mainly used for real-time speed test.
cv2.imshow('Frame', new_frame)
# Press Q on keyboard to exit
if cv2.waitKey(1) & 0xFF == ord('q'):
break