def run_classifier(c_value, kernel='linear'):
# multiclass support is handled by one-vs-one scheme
classifier = svm.SVC(C=c_value, kernel=kernel)
classifier.fit(X_train, y_train)
test_predicted = classifier.predict(X_test)
train_predicted = classifier.predict(X_train)
assert len(test_predicted) == len(y_test)
test_confusion = get_confusion_matrix(test_predicted, y_test)
train_confusion = get_confusion_matrix(train_predicted, y_train)
plot_heatmap(test_confusion,
f'../confusions/test_c_{c_value}_kernel_{kernel}.png')
plot_heatmap(train_confusion,
f'../confusions/train_c_{c_value}_kernel_{kernel}.png')
test_precision = get_precision(test_confusion)
train_precision = get_precision(train_confusion)
# correct = sum(1 for x in filter(lambda x: x, list(map(lambda x: x[0] == x[1], results))))
# print(f'There are {correct} out of {len(X_test)} samples. This gives us an accuracy of {correct/len(X_test)}')
print(f'Finished iteration with c {c_value} and kernel {kernel}')
return test_precision, train_precision
def get_optimizer(params):
"""Get optimizer."""
learning_rate = learning_rate_schedule(params)
momentum = params['momentum']
if params['optimizer'].lower() == 'sgd':
logging.info('Use SGD optimizer')
optimizer = tf.keras.optimizers.SGD(learning_rate, momentum=momentum)
elif params['optimizer'].lower() == 'adam':
logging.info('Use Adam optimizer')
optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=momentum)
else:
raise ValueError('optimizers should be adam or sgd')
moving_average_decay = params['moving_average_decay']
if moving_average_decay:
# TODO(tanmingxing): potentially add dynamic_decay for new tfa release.
from tensorflow_addons import optimizers as tfa_optimizers # pylint: disable=g-import-not-at-top
optimizer = tfa_optimizers.MovingAverage(
optimizer, average_decay=moving_average_decay, dynamic_decay=True)
precision = utils.get_precision(params['strategy'], params['mixed_precision'])
if precision == 'mixed_float16' and params['loss_scale']:
optimizer = tf.keras.mixed_precision.experimental.LossScaleOptimizer(
optimizer,
loss_scale=tf.mixed_precision.experimental.DynamicLossScale(
params['loss_scale']))
return optimizer
def __init__(self,
model_name: Text,
ckpt_path: Text = None,
batch_size: int = 1,
only_network: bool = False,
model_params: Dict[Text, Any] = None):
"""Initialize the inference driver.
Args:
model_name: target model name, such as efficientdet-d0.
ckpt_path: checkpoint path, such as /tmp/efficientdet-d0/.
batch_size: batch size for inference.
only_network: only use the network without pre/post processing.
model_params: model parameters for overriding the config.
"""
super().__init__()
self.model_name = model_name
self.ckpt_path = ckpt_path
self.batch_size = batch_size
self.only_network = only_network
self.params = hparams_config.get_detection_config(model_name).as_dict()
if model_params:
self.params.update(model_params)
self.params.update(dict(is_training_bn=False))
self.label_map = self.params.get('label_map', None)
self._model = None
mixed_precision = self.params.get('mixed_precision', None)
precision = utils.get_precision(self.params.get('strategy', None),
mixed_precision)
policy = tf.keras.mixed_precision.experimental.Policy(precision)
tf.keras.mixed_precision.experimental.set_policy(policy)
def build_model(model_name: Text, inputs: tf.Tensor, **kwargs):
"""Build model for a given model name.
Args:
model_name: the name of the model.
inputs: an image tensor or a numpy array.
**kwargs: extra parameters for model builder.
Returns:
(cls_outputs, box_outputs): the outputs for class and box predictions.
Each is a dictionary with key as feature level and value as predictions.
"""
model_arch = det_model_fn.get_model_arch(model_name)
mixed_precision = kwargs.get('mixed_precision', None)
precision = utils.get_precision(kwargs.get('strategy', None),
mixed_precision)
cls_outputs, box_outputs = utils.build_model_with_precision(
precision, model_arch, inputs, False, model_name, **kwargs)
if mixed_precision:
# Post-processing has multiple places with hard-coded float32.
# TODO(tanmingxing): Remove them once post-process can adpat to dtypes.
cls_outputs = {
k: tf.cast(v, tf.float32)
for k, v in cls_outputs.items()
}
box_outputs = {
k: tf.cast(v, tf.float32)
for k, v in box_outputs.items()
}
return cls_outputs, box_outputs
def __init__(self, ckpt_path, debug, *args, **kwargs):
""" Initialize the inference driver.
Args:
ckpt_path: checkpoint path, such as /tmp/efficientdet-d0/.
debug: bool, if true, run in debug mode.
"""
super().__init__(*args, **kwargs)
params = copy.deepcopy(self.params)
config = hparams_config.get_efficientdet_config(self.model_name)
config.override(params)
precision = utils.get_precision(config.strategy,
config.mixed_precision)
policy = tf.keras.mixed_precision.Policy(precision)
tf.keras.mixed_precision.set_global_policy(policy)
self.model = efficientdet_keras.EfficientDetModel(config=config)
image_size = utils.parse_image_size(config.image_size)
self.model.build((self.batch_size, *image_size, 3))
util_keras.restore_ckpt(self.model,
ckpt_path,
config.moving_average_decay,
skip_mismatch=False)
self.debug = debug
if debug:
tf.config.run_functions_eagerly(debug)
def launch(symbol):
try:
precision = get_precision(symbol)
pip = get_pip(symbol)
price_precision = len(
'{0:.10f}'.format(pip).split('.')[1].split('1')[0]) + 1
to_general_log(symbol, 'Start monitoring')
while is_allowed(symbol):
ma8 = get_ma_value(symbol, '5M', 8)
ma21 = get_ma_value(symbol, '5M', 21)
current_candle_open = float(get_current_candle(symbol, '5M')[1])
current_candle_close = float(get_current_candle(symbol, '5M')[4])
# looking for BUY signal
if ma21 < ma8 < min(current_candle_open,
current_candle_close) and check_anchor_chart(
symbol, 'buy'):
# waiting trigger bar for BUY
while True:
current_candle = get_current_candle(symbol, '5M')
current_candle_close = float(current_candle[4])
ma8 = get_ma_value(symbol, '5M', 8)
if current_candle_close < ma8:
entrance_point = round(
last_bars_extremum(symbol, 5, 'buy') + (30 * pip),
price_precision)
place_pending_order(symbol, 'buy', entrance_point,
precision, price_precision)
break
# looking for SELL signal
if ma21 > ma8 > max(current_candle_open,
current_candle_close) and check_anchor_chart(
symbol, 'sell'):
# waiting trigger bar for SELL
while True:
current_candle = get_current_candle(symbol, '5M')
current_candle_open = float(current_candle[1])
ma8 = get_ma_value(symbol, '5M', 8)
if current_candle_open > ma8:
entrance_point = round(
last_bars_extremum(symbol, 5, 'sell') - (30 * pip),
price_precision)
place_pending_order(symbol, 'sell', entrance_point,
precision, price_precision)
break
time.sleep(1)
except:
to_err_log(symbol, traceback.format_exc())
def build_model(model_name: Text, inputs: tf.Tensor, **kwargs):
"""Build model for a given model name.
Args:
model_name: the name of the model.
inputs: an image tensor or a numpy array.
**kwargs: extra parameters for model builder.
Returns:
(cls_outputs, box_outputs): the outputs for class and box predictions.
Each is a dictionary with key as feature level and value as predictions.
"""
mixed_precision = kwargs.get('mixed_precision', None)
precision = utils.get_precision(kwargs.get('strategy', None), mixed_precision)
if kwargs.get('use_keras_model', None):
def model_arch(feats, model_name=None, **kwargs):
"""Construct a model arch for keras models."""
config = hparams_config.get_efficientdet_config(model_name)
config.override(kwargs)
model = efficientdet_keras.EfficientDetNet(config=config)
#l=model.layers[0] # efficientnet part
#print(l.name)
#layer_names=[]
#feats_out=l.predict(feats,steps=1) #predict
#for ml in l.layers:
#print(ml.name)
#layer_names.append(ml.name)
#save_feat_fig(feats_out)
#exit()
cls_out_list, box_out_list = model(feats, training=False)
# convert the list of model outputs to a dictionary with key=level.
assert len(cls_out_list) == config.max_level - config.min_level + 1
assert len(box_out_list) == config.max_level - config.min_level + 1
cls_outputs, box_outputs = {}, {}
for i in range(config.min_level, config.max_level + 1):
cls_outputs[i] = cls_out_list[i - config.min_level]
box_outputs[i] = box_out_list[i - config.min_level]
return cls_outputs, box_outputs
else:
model_arch = det_model_fn.get_model_arch(model_name)
cls_outputs, box_outputs = utils.build_model_with_precision(
precision, model_arch, inputs, False, model_name, **kwargs)
if mixed_precision:
# Post-processing has multiple places with hard-coded float32.
# TODO(tanmingxing): Remove them once post-process can adpat to dtypes.
cls_outputs = {k: tf.cast(v, tf.float32) for k, v in cls_outputs.items()}
box_outputs = {k: tf.cast(v, tf.float32) for k, v in box_outputs.items()}
return cls_outputs, box_outputs
def __init__(self,
model_name: Text,
ckpt_path: Text = None,
batch_size: int = 1,
min_score_thresh: float = None,
max_boxes_to_draw: float = None,
model_params: Dict[Text, Any] = None):
"""Initialize the inference driver.
Args:
model_name: target model name, such as efficientdet-d0.
ckpt_path: checkpoint path, such as /tmp/efficientdet-d0/.
batch_size: batch size for inference.
min_score_thresh: minimal score threshold for filtering predictions.
max_boxes_to_draw: the maximum number of boxes per image.
model_params: model parameters for overriding the config.
"""
super().__init__()
self.model_name = model_name
self.ckpt_path = ckpt_path
self.batch_size = batch_size
self.params = hparams_config.get_detection_config(model_name).as_dict()
if model_params:
self.params.update(model_params)
self.params.update(dict(is_training_bn=False))
self.label_map = self.params.get('label_map', None)
self.model = None
self.min_score_thresh = min_score_thresh
self.max_boxes_to_draw = max_boxes_to_draw
mixed_precision = self.params.get('mixed_precision', None)
precision = utils.get_precision(self.params.get('strategy', None),
mixed_precision)
policy = tf.keras.mixed_precision.experimental.Policy(precision)
tf.keras.mixed_precision.experimental.set_policy(policy)
def launch(symbol):
try:
precision = get_precision(symbol)
pip = get_pip(symbol)
price_precision = len(
'{0:.10f}'.format(pip).split('.')[1].split('1')[0]) + 1
to_general_log(symbol, 'Start monitoring')
while is_allowed(symbol):
# looking for BUY signal
if get_price_change_percent_difference(symbol) > 2:
place_pending_order(symbol, 'buy', precision, price_precision)
break
# looking for SELL signal
if get_price_change_percent_difference(symbol) > 2:
place_pending_order(symbol, 'sell', precision, price_precision)
break
time.sleep(1)
except:
to_err_log(symbol, traceback.format_exc())
def main(_):
# 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
# Parse image size in case it is in string format.
config.image_size = utils.parse_image_size(config.image_size)
if FLAGS.use_xla and FLAGS.strategy != 'tpu':
tf.config.optimizer.set_jit(True)
for gpu in tf.config.list_physical_devices('GPU'):
tf.config.experimental.set_memory_growth(gpu, True)
if FLAGS.debug:
tf.debugging.set_log_device_placement(True)
os.environ['TF_DETERMINISTIC_OPS'] = '1'
tf.random.set_seed(FLAGS.tf_random_seed)
logging.set_verbosity(logging.DEBUG)
if FLAGS.strategy == 'tpu':
tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
FLAGS.tpu, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)
tf.config.experimental_connect_to_cluster(tpu_cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(tpu_cluster_resolver)
ds_strategy = tf.distribute.TPUStrategy(tpu_cluster_resolver)
logging.info('All devices: %s', tf.config.list_logical_devices('TPU'))
elif FLAGS.strategy == 'gpus':
gpus = tf.config.list_physical_devices('GPU')
if FLAGS.batch_size % len(gpus):
raise ValueError(
'Batch size divide gpus number must be interger, but got %f' %
(FLAGS.batch_size / len(gpus)))
if platform.system() == 'Windows':
# Windows doesn't support nccl use HierarchicalCopyAllReduce instead
# TODO(fsx950223): investigate HierarchicalCopyAllReduce performance issue
cross_device_ops = tf.distribute.HierarchicalCopyAllReduce()
else:
cross_device_ops = None
ds_strategy = tf.distribute.MirroredStrategy(
cross_device_ops=cross_device_ops)
logging.info('All devices: %s', gpus)
else:
if tf.config.list_physical_devices('GPU'):
ds_strategy = tf.distribute.OneDeviceStrategy('device:GPU:0')
else:
ds_strategy = tf.distribute.OneDeviceStrategy('device:CPU:0')
steps_per_epoch = FLAGS.num_examples_per_epoch // FLAGS.batch_size
params = dict(profile=FLAGS.profile,
model_name=FLAGS.model_name,
steps_per_execution=FLAGS.steps_per_execution,
model_dir=FLAGS.model_dir,
steps_per_epoch=steps_per_epoch,
strategy=FLAGS.strategy,
batch_size=FLAGS.batch_size,
tf_random_seed=FLAGS.tf_random_seed,
debug=FLAGS.debug,
val_json_file=FLAGS.val_json_file,
eval_samples=FLAGS.eval_samples,
num_shards=ds_strategy.num_replicas_in_sync)
config.override(params, True)
# set mixed precision policy by keras api.
precision = utils.get_precision(config.strategy, config.mixed_precision)
policy = tf.keras.mixed_precision.Policy(precision)
tf.keras.mixed_precision.set_global_policy(policy)
def get_dataset(is_training, config):
file_pattern = (FLAGS.train_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,
debug=FLAGS.debug)(config.as_dict())
with ds_strategy.scope():
if config.model_optimizations:
tfmot.set_config(config.model_optimizations.as_dict())
if FLAGS.hub_module_url:
model = train_lib.EfficientDetNetTrainHub(
config=config, hub_module_url=FLAGS.hub_module_url)
else:
model = train_lib.EfficientDetNetTrain(config=config)
model = setup_model(model, config)
if FLAGS.debug:
tf.config.run_functions_eagerly(True)
if FLAGS.pretrained_ckpt and not FLAGS.hub_module_url:
ckpt_path = tf.train.latest_checkpoint(FLAGS.pretrained_ckpt)
util_keras.restore_ckpt(model,
ckpt_path,
config.moving_average_decay,
exclude_layers=['class_net'])
init_experimental(config)
if 'train' in FLAGS.mode:
val_dataset = get_dataset(False,
config) if 'eval' in FLAGS.mode else None
model.fit(
get_dataset(True, config),
epochs=config.num_epochs,
steps_per_epoch=steps_per_epoch,
callbacks=train_lib.get_callbacks(config.as_dict(),
val_dataset),
validation_data=val_dataset,
validation_steps=(FLAGS.eval_samples // FLAGS.batch_size))
else:
# Continuous eval.
for ckpt in tf.train.checkpoints_iterator(FLAGS.model_dir,
min_interval_secs=180):
logging.info('Starting to evaluate.')
# Terminate eval job when final checkpoint is reached.
try:
current_epoch = int(os.path.basename(ckpt).split('-')[1])
except IndexError:
current_epoch = 0
val_dataset = get_dataset(False, config)
logging.info('start loading model.')
model.load_weights(tf.train.latest_checkpoint(FLAGS.model_dir))
logging.info('finish loading model.')
coco_eval = train_lib.COCOCallback(val_dataset, 1)
coco_eval.set_model(model)
eval_results = coco_eval.on_epoch_end(current_epoch)
logging.info('eval results for %s: %s', ckpt, eval_results)
try:
utils.archive_ckpt(eval_results, eval_results['AP'], ckpt)
except tf.errors.NotFoundError:
# Checkpoint might be not already deleted by the time eval finished.
logging.info('Checkpoint %s no longer exists, skipping.',
ckpt)
if current_epoch >= config.num_epochs or not current_epoch:
logging.info('Eval epoch %d / %d', current_epoch,
config.num_epochs)
break
def _detection_loss(self, cls_outputs, box_outputs, labels, loss_vals):
"""Computes total detection loss.
Computes total detection loss including box and class loss from all levels.
Args:
cls_outputs: an OrderDict with keys representing levels and values
representing logits in [batch_size, height, width, num_anchors].
box_outputs: an OrderDict with keys representing levels and values
representing box regression targets in [batch_size, height, width,
num_anchors * 4].
labels: the dictionary that returned from dataloader that includes
groundtruth targets.
loss_vals: A dict of loss values.
Returns:
total_loss: an integer tensor representing total loss reducing from
class and box losses from all levels.
cls_loss: an integer tensor representing total class loss.
box_loss: an integer tensor representing total box regression loss.
box_iou_loss: an integer tensor representing total box iou loss.
"""
# Sum all positives in a batch for normalization and avoid zero
# num_positives_sum, which would lead to inf loss during training
precision = utils.get_precision(self.config.strategy,
self.config.mixed_precision)
dtype = precision.split('_')[-1]
num_positives_sum = tf.reduce_sum(labels['mean_num_positives']) + 1.0
positives_momentum = self.config.positives_momentum or 0
if positives_momentum > 0:
# normalize the num_positive_examples for training stability.
moving_normalizer_var = tf.Variable(
0.0,
name='moving_normalizer',
dtype=dtype,
synchronization=tf.VariableSynchronization.ON_READ,
trainable=False,
aggregation=tf.VariableAggregation.MEAN)
num_positives_sum = tf.keras.backend.moving_average_update(
moving_normalizer_var,
num_positives_sum,
momentum=self.config.positives_momentum)
elif positives_momentum < 0:
num_positives_sum = utils.cross_replica_mean(num_positives_sum)
num_positives_sum = tf.cast(num_positives_sum, dtype)
levels = range(len(cls_outputs))
cls_losses = []
box_losses = []
for level in levels:
# Onehot encoding for classification labels.
cls_targets_at_level = tf.one_hot(
labels['cls_targets_%d' % (level + self.config.min_level)],
self.config.num_classes,
dtype=dtype)
if self.config.data_format == 'channels_first':
bs, _, width, height, _ = cls_targets_at_level.get_shape(
).as_list()
cls_targets_at_level = tf.reshape(cls_targets_at_level,
[bs, -1, width, height])
else:
bs, width, height, _, _ = cls_targets_at_level.get_shape(
).as_list()
cls_targets_at_level = tf.reshape(cls_targets_at_level,
[bs, width, height, -1])
class_loss_layer = self.loss.get(FocalLoss.__name__, None)
if class_loss_layer:
cls_loss = class_loss_layer(
[num_positives_sum, cls_targets_at_level],
cls_outputs[level])
if self.config.data_format == 'channels_first':
cls_loss = tf.reshape(
cls_loss,
[bs, -1, width, height, self.config.num_classes])
else:
cls_loss = tf.reshape(
cls_loss,
[bs, width, height, -1, self.config.num_classes])
cls_loss *= tf.cast(
tf.expand_dims(
tf.not_equal(
labels['cls_targets_%d' %
(level + self.config.min_level)], -2), -1),
dtype)
cls_loss_sum = tf.clip_by_value(tf.reduce_sum(cls_loss), 0.0,
2.0)
cls_losses.append(tf.cast(cls_loss_sum, dtype))
if self.config.box_loss_weight and self.loss.get(
BoxLoss.__name__, None):
box_targets_at_level = (
labels['box_targets_%d' % (level + self.config.min_level)])
box_loss_layer = self.loss[BoxLoss.__name__]
box_losses.append(
box_loss_layer([num_positives_sum, box_targets_at_level],
box_outputs[level]))
if self.config.iou_loss_type:
box_outputs = tf.concat(
[tf.reshape(v, [-1, 4]) for v in box_outputs], axis=0)
box_targets = tf.concat([
tf.reshape(
labels['box_targets_%d' %
(level + self.config.min_level)], [-1, 4])
for level in levels
],
axis=0)
box_iou_loss_layer = self.loss[BoxIouLoss.__name__]
box_iou_loss = box_iou_loss_layer([num_positives_sum, box_targets],
box_outputs)
loss_vals['box_iou_loss'] = box_iou_loss
else:
box_iou_loss = 0
cls_loss = tf.add_n(cls_losses) if cls_losses else 0
box_loss = tf.add_n(box_losses) if box_losses else 0
total_loss = (cls_loss + self.config.box_loss_weight * box_loss +
self.config.iou_loss_weight * box_iou_loss)
loss_vals['det_loss'] = total_loss
loss_vals['cls_loss'] = cls_loss
loss_vals['box_loss'] = box_loss
return total_loss
def _model_fn(features, labels, mode, params, model, variable_filter_fn=None):
"""Model definition entry.
Args:
features: the input image tensor with shape [batch_size, height, width, 3].
The height and width are fixed and equal.
labels: the input labels in a dictionary. The labels include class targets
and box targets which are dense label maps. The labels are generated from
get_input_fn function in data/dataloader.py
mode: the mode of TPUEstimator including TRAIN and EVAL.
params: the dictionary defines hyperparameters of model. The default
settings are in default_hparams function in this file.
model: the model outputs class logits and box regression outputs.
variable_filter_fn: the filter function that takes trainable_variables and
returns the variable list after applying the filter rule.
Returns:
tpu_spec: the TPUEstimatorSpec to run training, evaluation, or prediction.
Raises:
RuntimeError: if both ckpt and backbone_ckpt are set.
"""
is_tpu = params['strategy'] == 'tpu'
if params['img_summary_steps']:
utils.image('input_image', features, is_tpu)
training_hooks = []
params['is_training_bn'] = (mode == tf.estimator.ModeKeys.TRAIN)
if params['use_keras_model']:
def model_fn(inputs):
model = efficientdet_keras.EfficientDetNet(
config=hparams_config.Config(params))
cls_out_list, box_out_list = model(inputs,
params['is_training_bn'])
cls_outputs, box_outputs = {}, {}
for i in range(params['min_level'], params['max_level'] + 1):
cls_outputs[i] = cls_out_list[i - params['min_level']]
box_outputs[i] = box_out_list[i - params['min_level']]
return cls_outputs, box_outputs
else:
model_fn = functools.partial(model,
config=hparams_config.Config(params))
precision = utils.get_precision(params['strategy'],
params['mixed_precision'])
cls_outputs, box_outputs = utils.build_model_with_precision(
precision, model_fn, features, params['is_training_bn'])
levels = cls_outputs.keys()
for level in levels:
cls_outputs[level] = tf.cast(cls_outputs[level], tf.float32)
box_outputs[level] = tf.cast(box_outputs[level], tf.float32)
# Set up training loss and learning rate.
update_learning_rate_schedule_parameters(params)
global_step = tf.train.get_or_create_global_step()
learning_rate = learning_rate_schedule(params, global_step)
# cls_loss and box_loss are for logging. only total_loss is optimized.
det_loss, cls_loss, box_loss = detection_loss(cls_outputs, box_outputs,
labels, params)
reg_l2loss = reg_l2_loss(params['weight_decay'])
total_loss = det_loss + reg_l2loss
if mode == tf.estimator.ModeKeys.TRAIN:
utils.scalar('lrn_rate', learning_rate, is_tpu)
utils.scalar('trainloss/cls_loss', cls_loss, is_tpu)
utils.scalar('trainloss/box_loss', box_loss, is_tpu)
utils.scalar('trainloss/det_loss', det_loss, is_tpu)
utils.scalar('trainloss/reg_l2_loss', reg_l2loss, is_tpu)
utils.scalar('trainloss/loss', total_loss, is_tpu)
train_epochs = tf.cast(global_step,
tf.float32) / params['steps_per_epoch']
utils.scalar('train_epochs', train_epochs, is_tpu)
moving_average_decay = params['moving_average_decay']
if moving_average_decay:
ema = tf.train.ExponentialMovingAverage(decay=moving_average_decay,
num_updates=global_step)
ema_vars = utils.get_ema_vars()
if mode == tf.estimator.ModeKeys.TRAIN:
if params['optimizer'].lower() == 'sgd':
optimizer = tf.train.MomentumOptimizer(learning_rate,
momentum=params['momentum'])
elif params['optimizer'].lower() == 'adam':
optimizer = tf.train.AdamOptimizer(learning_rate)
else:
raise ValueError('optimizers should be adam or sgd')
if is_tpu:
optimizer = tf.tpu.CrossShardOptimizer(optimizer)
# Batch norm requires update_ops to be added as a train_op dependency.
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
var_list = tf.trainable_variables()
if variable_filter_fn:
var_list = variable_filter_fn(var_list)
if params.get('clip_gradients_norm', None):
logging.info('clip gradients norm by %f',
params['clip_gradients_norm'])
grads_and_vars = optimizer.compute_gradients(total_loss, var_list)
with tf.name_scope('clip'):
grads = [gv[0] for gv in grads_and_vars]
tvars = [gv[1] for gv in grads_and_vars]
# First clip each variable's norm, then clip global norm.
clip_norm = abs(params['clip_gradients_norm'])
clipped_grads = [
tf.clip_by_norm(g, clip_norm) if g is not None else None
for g in grads
]
clipped_grads, _ = tf.clip_by_global_norm(
clipped_grads, clip_norm)
utils.scalar('gradient_norm',
tf.linalg.global_norm(clipped_grads), is_tpu)
grads_and_vars = list(zip(clipped_grads, tvars))
with tf.control_dependencies(update_ops):
train_op = optimizer.apply_gradients(grads_and_vars,
global_step)
else:
with tf.control_dependencies(update_ops):
train_op = optimizer.minimize(total_loss,
global_step,
var_list=var_list)
if moving_average_decay:
with tf.control_dependencies([train_op]):
train_op = ema.apply(ema_vars)
else:
train_op = None
eval_metrics = None
if mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(**kwargs):
"""Returns a dictionary that has the evaluation metrics."""
if params['nms_configs'].get('pyfunc', True):
detections_bs = []
nms_configs = params['nms_configs']
for index in range(kwargs['boxes'].shape[0]):
detections = tf.numpy_function(
functools.partial(nms_np.per_class_nms,
nms_configs=nms_configs),
[
kwargs['boxes'][index],
kwargs['scores'][index],
kwargs['classes'][index],
tf.slice(kwargs['image_ids'], [index], [1]),
tf.slice(kwargs['image_scales'], [index], [1]),
params['num_classes'],
nms_configs['max_output_size'],
], tf.float32)
detections_bs.append(detections)
detections_bs = postprocess.transform_detections(
tf.stack(detections_bs))
else:
# These two branches should be equivalent, but currently they are not.
# TODO(tanmingxing): enable the non_pyfun path after bug fix.
nms_boxes, nms_scores, nms_classes, _ = postprocess.per_class_nms(
params, kwargs['boxes'], kwargs['scores'],
kwargs['classes'], kwargs['image_scales'])
img_ids = tf.cast(tf.expand_dims(kwargs['image_ids'], -1),
nms_scores.dtype)
detections_bs = [
img_ids * tf.ones_like(nms_scores),
nms_boxes[:, :, 1],
nms_boxes[:, :, 0],
nms_boxes[:, :, 3] - nms_boxes[:, :, 1],
nms_boxes[:, :, 2] - nms_boxes[:, :, 0],
nms_scores,
nms_classes,
]
detections_bs = tf.stack(detections_bs,
axis=-1,
name='detnections')
if params.get('testdev_dir', None):
logging.info('Eval testdev_dir %s', params['testdev_dir'])
eval_metric = coco_metric.EvaluationMetric(
testdev_dir=params['testdev_dir'])
coco_metrics = eval_metric.estimator_metric_fn(
detections_bs, tf.zeros([1]))
else:
logging.info('Eval val with groudtruths %s.',
params['val_json_file'])
eval_metric = coco_metric.EvaluationMetric(
filename=params['val_json_file'],
label_map=params['label_map'])
coco_metrics = eval_metric.estimator_metric_fn(
detections_bs, kwargs['groundtruth_data'])
# Add metrics to output.
cls_loss = tf.metrics.mean(kwargs['cls_loss_repeat'])
box_loss = tf.metrics.mean(kwargs['box_loss_repeat'])
output_metrics = {
'cls_loss': cls_loss,
'box_loss': box_loss,
}
output_metrics.update(coco_metrics)
return output_metrics
cls_loss_repeat = tf.reshape(
tf.tile(tf.expand_dims(cls_loss, 0), [
params['batch_size'],
]), [params['batch_size'], 1])
box_loss_repeat = tf.reshape(
tf.tile(tf.expand_dims(box_loss, 0), [
params['batch_size'],
]), [params['batch_size'], 1])
cls_outputs = postprocess.to_list(cls_outputs)
box_outputs = postprocess.to_list(box_outputs)
params['nms_configs']['max_nms_inputs'] = anchors.MAX_DETECTION_POINTS
boxes, scores, classes = postprocess.pre_nms(params, cls_outputs,
box_outputs)
metric_fn_inputs = {
'cls_loss_repeat': cls_loss_repeat,
'box_loss_repeat': box_loss_repeat,
'image_ids': labels['source_ids'],
'groundtruth_data': labels['groundtruth_data'],
'image_scales': labels['image_scales'],
'boxes': boxes,
'scores': scores,
'classes': classes,
}
eval_metrics = (metric_fn, metric_fn_inputs)
checkpoint = params.get('ckpt') or params.get('backbone_ckpt')
if checkpoint and mode == tf.estimator.ModeKeys.TRAIN:
# Initialize the model from an EfficientDet or backbone checkpoint.
if params.get('ckpt') and params.get('backbone_ckpt'):
raise RuntimeError(
'--backbone_ckpt and --checkpoint are mutually exclusive')
if params.get('backbone_ckpt'):
var_scope = params['backbone_name'] + '/'
if params['ckpt_var_scope'] is None:
# Use backbone name as default checkpoint scope.
ckpt_scope = params['backbone_name'] + '/'
else:
ckpt_scope = params['ckpt_var_scope'] + '/'
else:
# Load every var in the given checkpoint
var_scope = ckpt_scope = '/'
def scaffold_fn():
"""Loads pretrained model through scaffold function."""
logging.info('restore variables from %s', checkpoint)
var_map = utils.get_ckpt_var_map(
ckpt_path=checkpoint,
ckpt_scope=ckpt_scope,
var_scope=var_scope,
skip_mismatch=params['skip_mismatch'])
tf.train.init_from_checkpoint(checkpoint, var_map)
return tf.train.Scaffold()
elif mode == tf.estimator.ModeKeys.EVAL and moving_average_decay:
def scaffold_fn():
"""Load moving average variables for eval."""
logging.info('Load EMA vars with ema_decay=%f',
moving_average_decay)
restore_vars_dict = ema.variables_to_restore(ema_vars)
saver = tf.train.Saver(restore_vars_dict)
return tf.train.Scaffold(saver=saver)
else:
scaffold_fn = None
if is_tpu:
return tf.estimator.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
train_op=train_op,
eval_metrics=eval_metrics,
host_call=utils.get_tpu_host_call(global_step, params),
scaffold_fn=scaffold_fn,
training_hooks=training_hooks)
else:
# Profile every 1K steps.
if params.get('profile', False):
profile_hook = tf.estimator.ProfilerHook(
save_steps=1000,
output_dir=params['model_dir'],
show_memory=True)
training_hooks.append(profile_hook)
# Report memory allocation if OOM; it will slow down the running.
class OomReportingHook(tf.estimator.SessionRunHook):
def before_run(self, run_context):
return tf.estimator.SessionRunArgs(
fetches=[],
options=tf.RunOptions(
report_tensor_allocations_upon_oom=True))
training_hooks.append(OomReportingHook())
logging_hook = tf.estimator.LoggingTensorHook(
{
'step': global_step,
'det_loss': det_loss,
'cls_loss': cls_loss,
'box_loss': box_loss,
},
every_n_iter=params.get('iterations_per_loop', 100),
)
training_hooks.append(logging_hook)
eval_metric_ops = (eval_metrics[0](
**eval_metrics[1]) if eval_metrics else None)
return tf.estimator.EstimatorSpec(
mode=mode,
loss=total_loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops,
scaffold=scaffold_fn() if scaffold_fn else None,
training_hooks=training_hooks)
def _model_fn(features, labels, mode, params, model, variable_filter_fn=None):
"""Model definition entry.
Args:
features: the input image tensor with shape [batch_size, height, width, 3].
The height and width are fixed and equal.
labels: the input labels in a dictionary. The labels include class targets
and box targets which are dense label maps. The labels are generated from
get_input_fn function in data/dataloader.py
mode: the mode of TPUEstimator including TRAIN, EVAL, and PREDICT.
params: the dictionary defines hyperparameters of model. The default
settings are in default_hparams function in this file.
model: the model outputs class logits and box regression outputs.
variable_filter_fn: the filter function that takes trainable_variables and
returns the variable list after applying the filter rule.
Returns:
tpu_spec: the TPUEstimatorSpec to run training, evaluation, or prediction.
Raises:
RuntimeError: if both ckpt and backbone_ckpt are set.
"""
utils.image('input_image', features)
training_hooks = []
if params['data_format'] == 'channels_first':
features = tf.transpose(features, [0, 3, 1, 2])
def _model_outputs(inputs):
# Convert params (dict) to Config for easier access.
return model(inputs, config=hparams_config.Config(params))
precision = utils.get_precision(params['strategy'],
params['mixed_precision'])
cls_outputs, box_outputs = utils.build_model_with_precision(
precision, _model_outputs, features, params['is_training_bn'])
levels = cls_outputs.keys()
for level in levels:
cls_outputs[level] = tf.cast(cls_outputs[level], tf.float32)
box_outputs[level] = tf.cast(box_outputs[level], tf.float32)
# First check if it is in PREDICT mode.
if mode == tf.estimator.ModeKeys.PREDICT:
predictions = {
'image': features,
}
for level in levels:
predictions['cls_outputs_%d' % level] = cls_outputs[level]
predictions['box_outputs_%d' % level] = box_outputs[level]
return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
# Set up training loss and learning rate.
update_learning_rate_schedule_parameters(params)
global_step = tf.train.get_or_create_global_step()
learning_rate = learning_rate_schedule(params, global_step)
# cls_loss and box_loss are for logging. only total_loss is optimized.
det_loss, cls_loss, box_loss, box_iou_loss = detection_loss(
cls_outputs, box_outputs, labels, params)
reg_l2loss = reg_l2_loss(params['weight_decay'])
total_loss = det_loss + reg_l2loss
if mode == tf.estimator.ModeKeys.TRAIN:
utils.scalar('lrn_rate', learning_rate)
utils.scalar('trainloss/cls_loss', cls_loss)
utils.scalar('trainloss/box_loss', box_loss)
utils.scalar('trainloss/det_loss', det_loss)
utils.scalar('trainloss/reg_l2_loss', reg_l2loss)
utils.scalar('trainloss/loss', total_loss)
if params['iou_loss_type']:
utils.scalar('trainloss/box_iou_loss', box_iou_loss)
moving_average_decay = params['moving_average_decay']
if moving_average_decay:
ema = tf.train.ExponentialMovingAverage(decay=moving_average_decay,
num_updates=global_step)
ema_vars = utils.get_ema_vars()
if params['strategy'] == 'horovod':
import horovod.tensorflow as hvd # pylint: disable=g-import-not-at-top
learning_rate = learning_rate * hvd.size()
if mode == tf.estimator.ModeKeys.TRAIN:
if params['optimizer'].lower() == 'sgd':
optimizer = tf.train.MomentumOptimizer(learning_rate,
momentum=params['momentum'])
elif params['optimizer'].lower() == 'adam':
optimizer = tf.train.AdamOptimizer(learning_rate)
else:
raise ValueError('optimizers should be adam or sgd')
if params['strategy'] == 'tpu':
optimizer = tf.tpu.CrossShardOptimizer(optimizer)
elif params['strategy'] == 'horovod':
optimizer = hvd.DistributedOptimizer(optimizer)
training_hooks = [hvd.BroadcastGlobalVariablesHook(0)]
# Batch norm requires update_ops to be added as a train_op dependency.
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
var_list = tf.trainable_variables()
if variable_filter_fn:
var_list = variable_filter_fn(var_list)
if params.get('clip_gradients_norm', 0) > 0:
logging.info('clip gradients norm by %f',
params['clip_gradients_norm'])
grads_and_vars = optimizer.compute_gradients(total_loss, var_list)
with tf.name_scope('clip'):
grads = [gv[0] for gv in grads_and_vars]
tvars = [gv[1] for gv in grads_and_vars]
clipped_grads, gnorm = tf.clip_by_global_norm(
grads, params['clip_gradients_norm'])
utils.scalar('gnorm', gnorm)
grads_and_vars = list(zip(clipped_grads, tvars))
with tf.control_dependencies(update_ops):
train_op = optimizer.apply_gradients(grads_and_vars,
global_step)
else:
with tf.control_dependencies(update_ops):
train_op = optimizer.minimize(total_loss,
global_step,
var_list=var_list)
if moving_average_decay:
with tf.control_dependencies([train_op]):
train_op = ema.apply(ema_vars)
else:
train_op = None
eval_metrics = None
if mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(**kwargs):
"""Returns a dictionary that has the evaluation metrics."""
batch_size = params['batch_size']
if params['strategy'] == 'tpu':
batch_size = params['batch_size'] * params['num_shards']
eval_anchors = anchors.Anchors(params['min_level'],
params['max_level'],
params['num_scales'],
params['aspect_ratios'],
params['anchor_scale'],
params['image_size'])
anchor_labeler = anchors.AnchorLabeler(eval_anchors,
params['num_classes'])
cls_loss = tf.metrics.mean(kwargs['cls_loss_repeat'])
box_loss = tf.metrics.mean(kwargs['box_loss_repeat'])
if params.get('testdev_dir', None):
logging.info('Eval testdev_dir %s', params['testdev_dir'])
coco_metrics = coco_metric_fn(
batch_size,
anchor_labeler,
params['val_json_file'],
testdev_dir=params['testdev_dir'],
nms_configs=params['nms_configs'],
**kwargs)
else:
logging.info('Eval val with groudtruths %s.',
params['val_json_file'])
coco_metrics = coco_metric_fn(
batch_size,
anchor_labeler,
params['val_json_file'],
nms_configs=params['nms_configs'],
**kwargs)
# Add metrics to output.
output_metrics = {
'cls_loss': cls_loss,
'box_loss': box_loss,
}
output_metrics.update(coco_metrics)
return output_metrics
cls_loss_repeat = tf.reshape(
tf.tile(tf.expand_dims(cls_loss, 0), [
params['batch_size'],
]), [params['batch_size'], 1])
box_loss_repeat = tf.reshape(
tf.tile(tf.expand_dims(box_loss, 0), [
params['batch_size'],
]), [params['batch_size'], 1])
metric_fn_inputs = {
'cls_loss_repeat': cls_loss_repeat,
'box_loss_repeat': box_loss_repeat,
'source_ids': labels['source_ids'],
'groundtruth_data': labels['groundtruth_data'],
'image_scales': labels['image_scales'],
}
add_metric_fn_inputs(params, cls_outputs, box_outputs,
metric_fn_inputs)
eval_metrics = (metric_fn, metric_fn_inputs)
checkpoint = params.get('ckpt') or params.get('backbone_ckpt')
if checkpoint and mode == tf.estimator.ModeKeys.TRAIN:
# Initialize the model from an EfficientDet or backbone checkpoint.
if params.get('ckpt') and params.get('backbone_ckpt'):
raise RuntimeError(
'--backbone_ckpt and --checkpoint are mutually exclusive')
if params.get('backbone_ckpt'):
var_scope = params['backbone_name'] + '/'
if params['ckpt_var_scope'] is None:
# Use backbone name as default checkpoint scope.
ckpt_scope = params['backbone_name'] + '/'
else:
ckpt_scope = params['ckpt_var_scope'] + '/'
else:
# Load every var in the given checkpoint
var_scope = ckpt_scope = '/'
def scaffold_fn():
"""Loads pretrained model through scaffold function."""
logging.info('restore variables from %s', checkpoint)
var_map = utils.get_ckpt_var_map(ckpt_path=checkpoint,
ckpt_scope=ckpt_scope,
var_scope=var_scope,
var_exclude_expr=params.get(
'var_exclude_expr', None))
tf.train.init_from_checkpoint(checkpoint, var_map)
return tf.train.Scaffold()
elif mode == tf.estimator.ModeKeys.EVAL and moving_average_decay:
def scaffold_fn():
"""Load moving average variables for eval."""
logging.info('Load EMA vars with ema_decay=%f',
moving_average_decay)
restore_vars_dict = ema.variables_to_restore(ema_vars)
saver = tf.train.Saver(restore_vars_dict)
return tf.train.Scaffold(saver=saver)
else:
scaffold_fn = None
if params['strategy'] != 'tpu':
# Profile every 1K steps.
profile_hook = tf.train.ProfilerHook(save_steps=1000,
output_dir=params['model_dir'])
training_hooks.append(profile_hook)
# Report memory allocation if OOM
class OomReportingHook(tf.estimator.SessionRunHook):
def before_run(self, run_context):
return tf.estimator.SessionRunArgs(
fetches=[],
options=tf.RunOptions(
report_tensor_allocations_upon_oom=True))
training_hooks.append(OomReportingHook())
return tf.estimator.tpu.TPUEstimatorSpec(mode=mode,
loss=total_loss,
train_op=train_op,
eval_metrics=eval_metrics,
host_call=utils.get_tpu_host_call(
global_step, params),
scaffold_fn=scaffold_fn,
training_hooks=training_hooks)
def main(_):
# 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
# Parse image size in case it is in string format.
config.image_size = utils.parse_image_size(config.image_size)
if FLAGS.use_xla and FLAGS.strategy != 'tpu':
tf.config.optimizer.set_jit(True)
for gpu in tf.config.list_physical_devices('GPU'):
tf.config.experimental.set_memory_growth(gpu, True)
if FLAGS.debug:
tf.config.experimental_run_functions_eagerly(True)
tf.debugging.set_log_device_placement(True)
os.environ['TF_DETERMINISTIC_OPS'] = '1'
tf.random.set_seed(FLAGS.tf_random_seed)
logging.set_verbosity(logging.DEBUG)
if FLAGS.strategy == 'tpu':
tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
FLAGS.tpu, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)
tf.config.experimental_connect_to_cluster(tpu_cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(tpu_cluster_resolver)
ds_strategy = tf.distribute.TPUStrategy(tpu_cluster_resolver)
logging.info('All devices: %s', tf.config.list_logical_devices('TPU'))
elif FLAGS.strategy == 'gpus':
ds_strategy = tf.distribute.MirroredStrategy()
logging.info('All devices: %s', tf.config.list_physical_devices('GPU'))
else:
if tf.config.list_physical_devices('GPU'):
ds_strategy = tf.distribute.OneDeviceStrategy('device:GPU:0')
else:
ds_strategy = tf.distribute.OneDeviceStrategy('device:CPU:0')
steps_per_epoch = FLAGS.num_examples_per_epoch // FLAGS.batch_size
params = dict(profile=FLAGS.profile,
model_name=FLAGS.model_name,
iterations_per_loop=FLAGS.iterations_per_loop,
model_dir=FLAGS.model_dir,
steps_per_epoch=steps_per_epoch,
strategy=FLAGS.strategy,
batch_size=FLAGS.batch_size,
tf_random_seed=FLAGS.tf_random_seed,
debug=FLAGS.debug,
val_json_file=FLAGS.val_json_file,
eval_samples=FLAGS.eval_samples,
num_shards=ds_strategy.num_replicas_in_sync)
config.override(params, True)
# set mixed precision policy by keras api.
precision = utils.get_precision(config.strategy, config.mixed_precision)
policy = tf.keras.mixed_precision.experimental.Policy(precision)
tf.keras.mixed_precision.experimental.set_policy(policy)
def get_dataset(is_training, config):
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,
debug=FLAGS.debug)(config.as_dict())
with ds_strategy.scope():
if config.model_optimizations:
tfmot.set_config(config.model_optimizations.as_dict())
model = setup_model(config)
if FLAGS.pretrained_ckpt:
ckpt_path = tf.train.latest_checkpoint(FLAGS.pretrained_ckpt)
util_keras.restore_ckpt(model, ckpt_path)
init_experimental(config)
val_dataset = get_dataset(False, config).repeat()
model.fit(get_dataset(True, config),
epochs=config.num_epochs,
steps_per_epoch=steps_per_epoch,
callbacks=train_lib.get_callbacks(config.as_dict(),
val_dataset),
validation_data=val_dataset,
validation_steps=(FLAGS.eval_samples // FLAGS.batch_size))
model.save_weights(os.path.join(FLAGS.model_dir, 'ckpt-final'))
def main(_):
# 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
# Parse image size in case it is in string format.
config.image_size = utils.parse_image_size(config.image_size)
if FLAGS.use_xla and FLAGS.strategy != 'tpu':
tf.config.optimizer.set_jit(True)
for gpu in tf.config.list_physical_devices('GPU'):
tf.config.experimental.set_memory_growth(gpu, 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)
if FLAGS.strategy == 'tpu':
tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
FLAGS.tpu, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)
tf.config.experimental_connect_to_cluster(tpu_cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(tpu_cluster_resolver)
ds_strategy = tf.distribute.TPUStrategy(tpu_cluster_resolver)
logging.info('All devices: %s', tf.config.list_logical_devices('TPU'))
elif FLAGS.strategy == 'gpus':
ds_strategy = tf.distribute.MirroredStrategy()
logging.info('All devices: %s', tf.config.list_physical_devices('GPU'))
else:
if tf.config.list_physical_devices('GPU'):
ds_strategy = tf.distribute.OneDeviceStrategy('device:GPU:0')
else:
ds_strategy = tf.distribute.OneDeviceStrategy('device:CPU:0')
# Check data path
if FLAGS.mode in (
'train', 'train_and_eval') and FLAGS.training_file_pattern is None:
raise RuntimeError(
'You must specify --training_file_pattern for training.')
if FLAGS.mode in ('eval', 'train_and_eval'):
if FLAGS.validation_file_pattern is None:
raise RuntimeError('You must specify --validation_file_pattern '
'for evaluation.')
params = dict(config.as_dict(),
model_name=FLAGS.model_name,
iterations_per_loop=FLAGS.iterations_per_loop,
model_dir=FLAGS.model_dir,
num_examples_per_epoch=FLAGS.num_examples_per_epoch,
strategy=FLAGS.strategy,
batch_size=FLAGS.batch_size //
ds_strategy.num_replicas_in_sync,
num_shards=ds_strategy.num_replicas_in_sync,
val_json_file=FLAGS.val_json_file,
testdev_dir=FLAGS.testdev_dir,
mode=FLAGS.mode)
# set mixed precision policy by keras api.
precision = utils.get_precision(params['strategy'],
params['mixed_precision'])
policy = tf.keras.mixed_precision.experimental.Policy(precision)
tf.keras.mixed_precision.experimental.set_policy(policy)
def get_dataset(is_training, params):
file_pattern = (FLAGS.training_file_pattern
if is_training else FLAGS.validation_file_pattern)
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)(params)
with ds_strategy.scope():
model = train_lib.EfficientDetNetTrain(params['model_name'], config)
height, width = utils.parse_image_size(params['image_size'])
model.build((params['batch_size'], height, width, 3))
model.compile(
optimizer=train_lib.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':
train_lib.FocalLoss(params['alpha'],
params['gamma'],
label_smoothing=params['label_smoothing'],
reduction=tf.keras.losses.Reduction.NONE)
})
ckpt_path = tf.train.latest_checkpoint(FLAGS.model_dir)
if ckpt_path:
model.load_weights(ckpt_path)
model.freeze_vars(params['var_freeze_expr'])
model.fit(get_dataset(True, params=params),
steps_per_epoch=FLAGS.num_examples_per_epoch,
callbacks=train_lib.get_callbacks(params, FLAGS.profile),
validation_data=get_dataset(False, params=params),
validation_steps=FLAGS.eval_samples)
model.save_weights(os.path.join(FLAGS.model_dir, 'model'))
def main(_):
assert '.yaml' in FLAGS.config, 'Please provide path to yaml file.'
cfg = get_default_config()
cfg.merge_from_file(FLAGS.config)
cfg.freeze()
model_dir = FLAGS.model_dir
if not tf.io.gfile.exists(model_dir):
tf.io.gfile.makedirs(model_dir)
# init wandb
if cfg.WANDB.ENABLE:
wandb.tensorboard.patch(root_logdir=model_dir)
wandb.init(job_type='train',
group=cfg.WANDB.GROUP_NAME,
project=cfg.WANDB.PROJECT_NAME,
sync_tensorboard=cfg.WANDB.TENSORBOARD,
mode=cfg.WANDB.MODE,
config=dict(cfg),
resume=True)
if FLAGS.debug:
tf.config.run_functions_eagerly(True)
tf.debugging.set_log_device_placement(True)
os.environ['TF_DETERMINISTIC_OPS'] = '1'
tf.random.set_seed(1111)
logging.set_verbosity(logging.DEBUG)
tf.debugging.experimental.enable_dump_debug_info(
model_dir,
tensor_debug_mode="FULL_HEALTH",
circular_buffer_size=-1)
strategy = utils.get_strategy(FLAGS.num_gpus)
# mixed precision
precision = utils.get_precision(FLAGS.mixed_precision)
policy = tf.keras.mixed_precision.Policy(precision)
tf.keras.mixed_precision.set_global_policy(policy)
def get_dataset(cfg, file_pattern, is_training):
"""Returns a tf.data.Dataset"""
return dataloader.InputReader(
cfg, is_training, FLAGS.use_tfrecord, FLAGS.mixed_precision)(
file_pattern,
cfg.TRAIN.BATCH_SIZE if is_training else cfg.TEST.BATCH_SIZE)
def load_model(model, cfg):
"""Compile model with loss function, model optimizers and metrics."""
opt_str = cfg.TRAIN.OPTIMIZER.lower()
if opt_str == 'sgd':
opt = tf.optimizers.SGD(learning_rate=cfg.TRAIN.WARMUP_LR,
momentum=cfg.TRAIN.MOMENTUM,
nesterov=True)
elif opt_str == 'adam':
opt = tf.optimizers.Adam(learning_rate=cfg.TRAIN.WARMUP_LR)
else:
raise NotImplementedError(f'{opt_str} not supported')
if FLAGS.mixed_precision:
opt = tf.keras.mixed_precision.LossScaleOptimizer(opt)
model.compile(
optimizer=opt,
loss=tf.keras.losses.SparseCategoricalCrossentropy(),
metrics=[
tf.keras.metrics.SparseCategoricalAccuracy(name='acc'),
tf.keras.metrics.SparseTopKCategoricalAccuracy(
k=5, name='top_5_acc')
])
return model
# learning rate schedule
def lr_schedule(epoch, lr):
"""
Implements the learning rate schedule used in
https://arxiv.org/abs/2004.04730
"""
if epoch > cfg.TRAIN.WARMUP_EPOCHS:
new_lr = cfg.TRAIN.BASE_LR * (0.5 * (
tf.math.cos(tf.constant(math.pi) *
(epoch / cfg.TRAIN.EPOCHS)) + 1))
else:
new_lr = cfg.TRAIN.WARMUP_LR + (
epoch * (cfg.TRAIN.BASE_LR - cfg.TRAIN.WARMUP_LR) /
cfg.TRAIN.WARMUP_EPOCHS)
return new_lr
with strategy.scope():
model = X3D(cfg)
model = load_model(model, cfg)
# resume training from latest checkpoint, if available
current_epoch = 0
ckpt_path = tf.train.latest_checkpoint(model_dir)
if ckpt_path:
current_epoch = int(os.path.basename(ckpt_path).split('-')[1])
logging.info(
f'Found checkpoint {ckpt_path} at epoch {current_epoch}')
model.load_weights(ckpt_path)
elif FLAGS.pretrained_ckpt:
logging.info(
f'Loading model from pretrained weights at {FLAGS.pretrained_ckpt}'
)
if tf.io.gfile.isdir(FLAGS.pretrained_ckpt):
model.load_weights(
tf.train.latest_checkpoint(FLAGS.pretrained_ckpt))
else:
model.load_weights(FLAGS.pretrained_ckpt)
model.fit(
get_dataset(cfg, FLAGS.train_file_pattern, True),
verbose=1,
epochs=cfg.TRAIN.EPOCHS,
initial_epoch=current_epoch,
steps_per_epoch=cfg.TRAIN.DATASET_SIZE // cfg.TRAIN.BATCH_SIZE,
validation_data=get_dataset(cfg, FLAGS.val_file_pattern, False)
if FLAGS.val_file_pattern else None,
callbacks=utils.get_callbacks(cfg, lr_schedule, FLAGS))
def main(_):
# 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
# Parse image size in case it is in string format.
config.image_size = utils.parse_image_size(config.image_size)
if FLAGS.use_xla and FLAGS.strategy != 'tpu':
tf.config.optimizer.set_jit(True)
for gpu in tf.config.list_physical_devices('GPU'):
tf.config.experimental.set_memory_growth(gpu, 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)
if FLAGS.strategy == 'tpu':
tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
FLAGS.tpu, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)
tf.config.experimental_connect_to_cluster(tpu_cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(tpu_cluster_resolver)
ds_strategy = tf.distribute.TPUStrategy(tpu_cluster_resolver)
logging.info('All devices: %s', tf.config.list_logical_devices('TPU'))
elif FLAGS.strategy == 'gpus':
ds_strategy = tf.distribute.MirroredStrategy()
logging.info('All devices: %s', tf.config.list_physical_devices('GPU'))
else:
if tf.config.list_physical_devices('GPU'):
ds_strategy = tf.distribute.OneDeviceStrategy('device:GPU:0')
else:
ds_strategy = tf.distribute.OneDeviceStrategy('device:CPU:0')
steps_per_epoch = FLAGS.num_examples_per_epoch // FLAGS.batch_size
params = dict(config.as_dict(),
profile=FLAGS.profile,
model_name=FLAGS.model_name,
iterations_per_loop=FLAGS.iterations_per_loop,
model_dir=FLAGS.model_dir,
steps_per_epoch=steps_per_epoch,
strategy=FLAGS.strategy,
batch_size=FLAGS.batch_size,
num_shards=ds_strategy.num_replicas_in_sync)
# set mixed precision policy by keras api.
precision = utils.get_precision(params['strategy'],
params['mixed_precision'])
policy = tf.keras.mixed_precision.experimental.Policy(precision)
tf.keras.mixed_precision.experimental.set_policy(policy)
def get_dataset(is_training, params):
file_pattern = (FLAGS.training_file_pattern
if is_training else FLAGS.validation_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)(params)
with ds_strategy.scope():
model = train_lib.EfficientDetNetTrain(params['model_name'], config)
model.compile(
optimizer=train_lib.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':
train_lib.FocalLoss(params['alpha'],
params['gamma'],
label_smoothing=params['label_smoothing'],
reduction=tf.keras.losses.Reduction.NONE),
'seg_loss':
tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True, reduction=tf.keras.losses.Reduction.NONE)
})
if FLAGS.pretrained_ckpt:
ckpt_path = tf.train.latest_checkpoint(FLAGS.pretrained_ckpt)
util_keras.restore_ckpt(model, ckpt_path,
params['moving_average_decay'])
tf.io.gfile.makedirs(FLAGS.model_dir)
if params['model_optimizations']:
model_optimization.set_config(params['model_optimizations'])
model.build((FLAGS.batch_size, *config.image_size, 3))
model.fit(get_dataset(True, params=params),
epochs=params['num_epochs'],
steps_per_epoch=steps_per_epoch,
callbacks=train_lib.get_callbacks(params),
validation_data=get_dataset(False, params=params).repeat(),
validation_steps=(FLAGS.eval_samples // FLAGS.batch_size))
model.save_weights(os.path.join(FLAGS.model_dir, 'ckpt-final'))
def _model_fn(features, labels, mode, params, model, variable_filter_fn=None):
"""Model definition entry.
Args:
features: the input image tensor with shape [batch_size, height, width, 3].
The height and width are fixed and equal.
labels: the input labels in a dictionary. The labels include class targets
and box targets which are dense label maps. The labels are generated from
get_input_fn function in data/dataloader.py
mode: the mode of TPUEstimator including TRAIN, EVAL, and PREDICT.
params: the dictionary defines hyperparameters of model. The default
settings are in default_hparams function in this file.
model: the model outputs class logits and box regression outputs.
variable_filter_fn: the filter function that takes trainable_variables and
returns the variable list after applying the filter rule.
Returns:
tpu_spec: the TPUEstimatorSpec to run training, evaluation, or prediction.
Raises:
RuntimeError: if both ckpt and backbone_ckpt are set.
"""
utils.image('input_image', features)
training_hooks = []
params['is_training_bn'] = (mode == tf.estimator.ModeKeys.TRAIN)
if params['use_keras_model']:
def model_fn(inputs):
model = efficientdet_keras.EfficientDetNet(
config=hparams_config.Config(params))
cls_out_list, box_out_list = model(inputs,
params['is_training_bn'])
cls_outputs, box_outputs = {}, {}
for i in range(params['min_level'], params['max_level'] + 1):
cls_outputs[i] = cls_out_list[i - params['min_level']]
box_outputs[i] = box_out_list[i - params['min_level']]
return cls_outputs, box_outputs
else:
model_fn = functools.partial(model,
config=hparams_config.Config(params))
precision = utils.get_precision(params['strategy'],
params['mixed_precision'])
cls_outputs, box_outputs = utils.build_model_with_precision(
precision, model_fn, features, params['is_training_bn'])
levels = cls_outputs.keys()
for level in levels:
cls_outputs[level] = tf.cast(cls_outputs[level], tf.float32)
box_outputs[level] = tf.cast(box_outputs[level], tf.float32)
# First check if it is in PREDICT mode.
if mode == tf.estimator.ModeKeys.PREDICT:
predictions = {
'image': features,
}
for level in levels:
predictions['cls_outputs_%d' % level] = cls_outputs[level]
predictions['box_outputs_%d' % level] = box_outputs[level]
return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
# Set up training loss and learning rate.
update_learning_rate_schedule_parameters(params)
global_step = tf.train.get_or_create_global_step()
learning_rate = learning_rate_schedule(params, global_step)
# cls_loss and box_loss are for logging. only total_loss is optimized.
det_loss, cls_loss, box_loss, box_iou_loss = detection_loss(
cls_outputs, box_outputs, labels, params)
reg_l2loss = reg_l2_loss(params['weight_decay'])
total_loss = det_loss + reg_l2loss
if mode == tf.estimator.ModeKeys.TRAIN:
utils.scalar('lrn_rate', learning_rate)
utils.scalar('trainloss/cls_loss', cls_loss)
utils.scalar('trainloss/box_loss', box_loss)
utils.scalar('trainloss/det_loss', det_loss)
utils.scalar('trainloss/reg_l2_loss', reg_l2loss)
utils.scalar('trainloss/loss', total_loss)
if params['iou_loss_type']:
utils.scalar('trainloss/box_iou_loss', box_iou_loss)
train_epochs = tf.cast(global_step,
tf.float32) / params['steps_per_epoch']
utils.scalar('train_epochs', train_epochs)
moving_average_decay = params['moving_average_decay']
if moving_average_decay:
ema = tf.train.ExponentialMovingAverage(decay=moving_average_decay,
num_updates=global_step)
ema_vars = utils.get_ema_vars()
if mode == tf.estimator.ModeKeys.TRAIN:
if params['optimizer'].lower() == 'sgd':
optimizer = tf.train.MomentumOptimizer(learning_rate,
momentum=params['momentum'])
elif params['optimizer'].lower() == 'adam':
optimizer = tf.train.AdamOptimizer(learning_rate)
else:
raise ValueError('optimizers should be adam or sgd')
if params['strategy'] == 'tpu':
optimizer = tf.tpu.CrossShardOptimizer(optimizer)
if params['gradient_checkpointing']:
from third_party.grad_checkpoint \
import memory_saving_gradients # pylint: disable=g-import-not-at-top
from tensorflow.python.ops \
import gradients # pylint: disable=g-import-not-at-top
# monkey patch tf.gradients to point to our custom version,
# with automatic checkpoint selection
def gradients_(ys, xs, grad_ys=None, **kwargs):
return memory_saving_gradients.gradients(
ys,
xs,
grad_ys,
checkpoints=params['gradient_checkpointing_list'],
**kwargs)
gradients.__dict__["gradients"] = gradients_
# Batch norm requires update_ops to be added as a train_op dependency.
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
var_list = tf.trainable_variables()
if variable_filter_fn:
var_list = variable_filter_fn(var_list)
if params.get('clip_gradients_norm', None):
logging.info('clip gradients norm by %f',
params['clip_gradients_norm'])
grads_and_vars = optimizer.compute_gradients(total_loss, var_list)
with tf.name_scope('clip'):
grads = [gv[0] for gv in grads_and_vars]
tvars = [gv[1] for gv in grads_and_vars]
# First clip each variable's norm, then clip global norm.
clip_norm = abs(params['clip_gradients_norm'])
clipped_grads = [tf.clip_by_norm(g, clip_norm) for g in grads]
clipped_grads, _ = tf.clip_by_global_norm(
clipped_grads, clip_norm)
utils.scalar('gradient_norm',
tf.linalg.global_norm(clipped_grads))
grads_and_vars = list(zip(clipped_grads, tvars))
with tf.control_dependencies(update_ops):
train_op = optimizer.apply_gradients(grads_and_vars,
global_step)
else:
with tf.control_dependencies(update_ops):
train_op = optimizer.minimize(total_loss,
global_step,
var_list=var_list)
if moving_average_decay:
with tf.control_dependencies([train_op]):
train_op = ema.apply(ema_vars)
else:
train_op = None
eval_metrics = None
if mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(**kwargs):
"""Returns a dictionary that has the evaluation metrics."""
if params['nms_configs'].get('pyfunc', True):
detections_bs = []
for index in range(kwargs['boxes'].shape[0]):
nms_configs = params['nms_configs']
detections = tf.numpy_function(
functools.partial(nms_np.per_class_nms,
nms_configs=nms_configs),
[
kwargs['boxes'][index],
kwargs['scores'][index],
kwargs['classes'][index],
tf.slice(kwargs['image_ids'], [index], [1]),
tf.slice(kwargs['image_scales'], [index], [1]),
params['num_classes'],
nms_configs['max_output_size'],
], tf.float32)
detections_bs.append(detections)
detections_bs = postprocess.transform_detections(
tf.stack(detections_bs))
else:
# These two branches should be equivalent, but currently they are not.
# TODO(tanmingxing): enable the non_pyfun path after bug fix.
nms_boxes, nms_scores, nms_classes, _ = postprocess.per_class_nms(
params, kwargs['boxes'], kwargs['scores'],
kwargs['classes'], kwargs['image_scales'])
img_ids = tf.cast(tf.expand_dims(kwargs['image_ids'], -1),
nms_scores.dtype)
detections_bs = [
img_ids * tf.ones_like(nms_scores),
nms_boxes[:, :, 1],
nms_boxes[:, :, 0],
nms_boxes[:, :, 3] - nms_boxes[:, :, 1],
nms_boxes[:, :, 2] - nms_boxes[:, :, 0],
nms_scores,
nms_classes,
]
detections_bs = tf.stack(detections_bs,
axis=-1,
name='detnections')
if params.get('testdev_dir', None):
logging.info('Eval testdev_dir %s', params['testdev_dir'])
eval_metric = coco_metric.EvaluationMetric(
testdev_dir=params['testdev_dir'])
coco_metrics = eval_metric.estimator_metric_fn(
detections_bs, tf.zeros([1]))
else:
logging.info('Eval val with groudtruths %s.',
params['val_json_file'])
eval_metric = coco_metric.EvaluationMetric(
filename=params['val_json_file'])
coco_metrics = eval_metric.estimator_metric_fn(
detections_bs, kwargs['groundtruth_data'],
params['label_map'])
# Add metrics to output.
cls_loss = tf.metrics.mean(kwargs['cls_loss_repeat'])
box_loss = tf.metrics.mean(kwargs['box_loss_repeat'])
output_metrics = {
'cls_loss': cls_loss,
'box_loss': box_loss,
}
output_metrics.update(coco_metrics)
return output_metrics
cls_loss_repeat = tf.reshape(
tf.tile(tf.expand_dims(cls_loss, 0), [
params['batch_size'],
]), [params['batch_size'], 1])
box_loss_repeat = tf.reshape(
tf.tile(tf.expand_dims(box_loss, 0), [
params['batch_size'],
]), [params['batch_size'], 1])
cls_outputs = postprocess.to_list(cls_outputs)
box_outputs = postprocess.to_list(box_outputs)
params['nms_configs']['max_nms_inputs'] = anchors.MAX_DETECTION_POINTS
boxes, scores, classes = postprocess.pre_nms(params, cls_outputs,
box_outputs)
metric_fn_inputs = {
'cls_loss_repeat': cls_loss_repeat,
'box_loss_repeat': box_loss_repeat,
'image_ids': labels['source_ids'],
'groundtruth_data': labels['groundtruth_data'],
'image_scales': labels['image_scales'],
'boxes': boxes,
'scores': scores,
'classes': classes,
}
eval_metrics = (metric_fn, metric_fn_inputs)
checkpoint = params.get('ckpt') or params.get('backbone_ckpt')
if checkpoint and mode == tf.estimator.ModeKeys.TRAIN:
# Initialize the model from an EfficientDet or backbone checkpoint.
if params.get('ckpt') and params.get('backbone_ckpt'):
raise RuntimeError(
'--backbone_ckpt and --checkpoint are mutually exclusive')
if params.get('backbone_ckpt'):
var_scope = params['backbone_name'] + '/'
if params['ckpt_var_scope'] is None:
# Use backbone name as default checkpoint scope.
ckpt_scope = params['backbone_name'] + '/'
else:
ckpt_scope = params['ckpt_var_scope'] + '/'
else:
# Load every var in the given checkpoint
var_scope = ckpt_scope = '/'
def scaffold_fn():
"""Loads pretrained model through scaffold function."""
logging.info('restore variables from %s', checkpoint)
var_map = utils.get_ckpt_var_map(
ckpt_path=checkpoint,
ckpt_scope=ckpt_scope,
var_scope=var_scope,
skip_mismatch=params['skip_mismatch'])
tf.train.init_from_checkpoint(checkpoint, var_map)
return tf.train.Scaffold()
elif mode == tf.estimator.ModeKeys.EVAL and moving_average_decay:
def scaffold_fn():
"""Load moving average variables for eval."""
logging.info('Load EMA vars with ema_decay=%f',
moving_average_decay)
restore_vars_dict = ema.variables_to_restore(ema_vars)
saver = tf.train.Saver(restore_vars_dict)
return tf.train.Scaffold(saver=saver)
else:
scaffold_fn = None
if params['strategy'] != 'tpu':
# Profile every 1K steps.
if params.get('profile', False):
profile_hook = tf.estimator.ProfilerHook(
save_steps=1000,
output_dir=params['model_dir'],
show_memory=True)
training_hooks.append(profile_hook)
# Report memory allocation if OOM
class OomReportingHook(tf.estimator.SessionRunHook):
def before_run(self, run_context):
return tf.estimator.SessionRunArgs(
fetches=[],
options=tf.RunOptions(
report_tensor_allocations_upon_oom=True))
training_hooks.append(OomReportingHook())
logging_hook = tf.estimator.LoggingTensorHook(
{
'step': global_step,
'det_loss': det_loss,
'cls_loss': cls_loss,
'box_loss': box_loss,
},
every_n_iter=params.get('iterations_per_loop', 100),
)
training_hooks.append(logging_hook)
if params["nvgpu_logging"]:
try:
from third_party import nvgpu # pylint: disable=g-import-not-at-top
from functools import reduce # pylint: disable=g-import-not-at-top
def get_nested_value(d, path):
return reduce(dict.get, path, d)
def nvgpu_gpu_info(inp):
inp = inp.decode("utf-8")
inp = inp.split(",")
inp = [x.strip() for x in inp]
value = get_nested_value(nvgpu.gpu_info(), inp)
return np.str(value)
def commonsize(inp):
const_sizes = {
'B': 1,
'KB': 1e3,
'MB': 1e6,
'GB': 1e9,
'TB': 1e12,
'PB': 1e15,
'KiB': 1024,
'MiB': 1048576,
'GiB': 1073741824
}
inp = inp.split(" ")
# convert all to MiB
if inp[1] != 'MiB':
inp_ = float(
inp[0]) * (const_sizes[inp[1]] / 1048576.0)
else:
inp_ = float(inp[0])
return inp_
def formatter_log(tensors):
"""Format the output."""
mem_used = tensors["memory used"].decode("utf-8")
mem_total = tensors["memory total"].decode("utf-8")
mem_util = commonsize(mem_used) / commonsize(mem_total)
logstring = "GPU memory used: {} = {:.1%} of total GPU memory: {}".format(
mem_used, mem_util, mem_total)
return logstring
mem_used = tf.py_func(nvgpu_gpu_info,
['gpu, fb_memory_usage, used'],
[tf.string])[0]
mem_total = tf.py_func(nvgpu_gpu_info,
['gpu, fb_memory_usage, total'],
[tf.string])[0]
logging_hook3 = tf.estimator.LoggingTensorHook(
tensors={
"memory used": mem_used,
"memory total": mem_total,
},
every_n_iter=params.get('iterations_per_loop', 100),
formatter=formatter_log,
)
training_hooks.append(logging_hook3)
except:
logging.error("nvgpu error: nvidia-smi format not recognized")
if params['strategy'] == 'tpu':
return tf.estimator.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
train_op=train_op,
eval_metrics=eval_metrics,
host_call=utils.get_tpu_host_call(global_step, params),
scaffold_fn=scaffold_fn,
training_hooks=training_hooks)
else:
eval_metric_ops = eval_metrics[0](
**eval_metrics[1]) if eval_metrics else None
utils.get_tpu_host_call(global_step, params)
return tf.estimator.EstimatorSpec(mode=mode,
loss=total_loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops,
scaffold=scaffold_fn(),
training_hooks=training_hooks)
def dataset_parser(self, value, example_decoder, anchor_labeler, params):
"""Parse data to a fixed dimension input image and learning targets.
Args:
value: a single serialized tf.Example string.
example_decoder: TF example decoder.
anchor_labeler: anchor box labeler.
params: a dict of extra parameters.
Returns:
image: Image tensor that is preprocessed to have normalized value and
fixed dimension [image_height, image_width, 3]
cls_targets_dict: ordered dictionary with keys
[min_level, min_level+1, ..., max_level]. The values are tensor with
shape [height_l, width_l, num_anchors]. The height_l and width_l
represent the dimension of class logits at l-th level.
box_targets_dict: ordered dictionary with keys
[min_level, min_level+1, ..., max_level]. The values are tensor with
shape [height_l, width_l, num_anchors * 4]. The height_l and
width_l represent the dimension of bounding box regression output at
l-th level.
num_positives: Number of positive anchors in the image.
source_id: Source image id. Default value -1 if the source id is empty
in the groundtruth annotation.
image_scale: Scale of the processed image to the original image.
boxes: Groundtruth bounding box annotations. The box is represented in
[y1, x1, y2, x2] format. The tensor is padded with -1 to the fixed
dimension [self._max_instances_per_image, 4].
is_crowds: Groundtruth annotations to indicate if an annotation
represents a group of instances by value {0, 1}. The tensor is
padded with 0 to the fixed dimension [self._max_instances_per_image].
areas: Groundtruth areas annotations. The tensor is padded with -1
to the fixed dimension [self._max_instances_per_image].
classes: Groundtruth classes annotations. The tensor is padded with -1
to the fixed dimension [self._max_instances_per_image].
"""
with tf.name_scope('parser'):
data = example_decoder.decode(value)
source_id = data['source_id']
image = data['image']
boxes = data['groundtruth_boxes']
classes = data['groundtruth_classes']
classes = tf.reshape(tf.cast(classes, dtype=tf.float32), [-1, 1])
areas = data['groundtruth_area']
is_crowds = data['groundtruth_is_crowd']
image_masks = data.get('groundtruth_instance_masks', [])
classes = tf.reshape(tf.cast(classes, dtype=tf.float32), [-1, 1])
if self._is_training:
# Training time preprocessing.
if params['skip_crowd_during_training']:
indices = tf.where(
tf.logical_not(data['groundtruth_is_crowd']))
classes = tf.gather_nd(classes, indices)
boxes = tf.gather_nd(boxes, indices)
if params.get('grid_mask', None):
from aug import gridmask # pylint: disable=g-import-not-at-top
image, boxes = gridmask.gridmask(image, boxes)
if params.get('autoaugment_policy', None):
from aug import autoaugment # pylint: disable=g-import-not-at-top
if params['autoaugment_policy'] == 'randaug':
image, boxes = autoaugment.distort_image_with_randaugment(
image, boxes, num_layers=1, magnitude=15)
else:
image, boxes = autoaugment.distort_image_with_autoaugment(
image, boxes, params['autoaugment_policy'])
input_processor = DetectionInputProcessor(image,
params['image_size'],
boxes, classes)
input_processor.normalize_image()
if self._is_training:
if params['input_rand_hflip']:
input_processor.random_horizontal_flip()
input_processor.set_training_random_scale_factors(
params['jitter_min'], params['jitter_max'],
params.get('target_size', None))
else:
input_processor.set_scale_factors_to_output_size()
image = input_processor.resize_and_crop_image()
boxes, classes = input_processor.resize_and_crop_boxes()
# Assign anchors.
(cls_targets, box_targets,
num_positives) = anchor_labeler.label_anchors(boxes, classes)
source_id = tf.where(tf.equal(source_id, tf.constant('')), '-1',
source_id)
source_id = tf.strings.to_number(source_id)
# Pad groundtruth data for evaluation.
image_scale = input_processor.image_scale_to_original
boxes *= image_scale
is_crowds = tf.cast(is_crowds, dtype=tf.float32)
boxes = pad_to_fixed_size(boxes, -1,
[self._max_instances_per_image, 4])
is_crowds = pad_to_fixed_size(is_crowds, 0,
[self._max_instances_per_image, 1])
areas = pad_to_fixed_size(areas, -1,
[self._max_instances_per_image, 1])
classes = pad_to_fixed_size(classes, -1,
[self._max_instances_per_image, 1])
if params['mixed_precision']:
precision = utils.get_precision(params['strategy'],
params['mixed_precision'])
dtype = precision.split('_')[-1]
image = tf.cast(image, dtype=dtype)
box_targets = tf.nest.map_structure(
lambda box_target: tf.cast(box_target, dtype=dtype),
box_targets)
return (image, cls_targets, box_targets, num_positives, source_id,
image_scale, boxes, is_crowds, areas, classes, image_masks)
