def setup_model(config):
"""Build and compile model."""
model = train_lib.EfficientDetNetTrain(config=config)
model.build((None, *config.image_size, 3))
model.compile(
optimizer=train_lib.get_optimizer(config.as_dict()),
loss={
train_lib.BoxLoss.__name__:
train_lib.BoxLoss(config.delta,
reduction=tf.keras.losses.Reduction.NONE),
train_lib.BoxIouLoss.__name__:
train_lib.BoxIouLoss(config.iou_loss_type,
config.min_level,
config.max_level,
config.num_scales,
config.aspect_ratios,
config.anchor_scale,
config.image_size,
reduction=tf.keras.losses.Reduction.NONE),
train_lib.FocalLoss.__name__:
train_lib.FocalLoss(config.alpha,
config.gamma,
label_smoothing=config.label_smoothing,
reduction=tf.keras.losses.Reduction.NONE),
tf.keras.losses.SparseCategoricalCrossentropy.__name__:
tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True, reduction=tf.keras.losses.Reduction.NONE)
})
return model
def test_losses(self):
tf.random.set_seed(1111)
box_loss = train_lib.BoxLoss()
box_iou_loss = train_lib.BoxIouLoss(iou_loss_type='ciou',
min_level=3,
max_level=3,
num_scales=1,
aspect_ratios=[1.0],
anchor_scale=1.0,
image_size=32)
alpha = 0.25
gamma = 1.5
focal_loss_v2 = train_lib.FocalLoss(
alpha, gamma, reduction=tf.keras.losses.Reduction.NONE)
box_outputs = tf.random.normal([64, 4])
box_targets = tf.random.normal([64, 4])
num_positives = tf.constant(4.0)
self.assertEqual(
legacy_fn._box_loss(box_outputs, box_targets, num_positives),
box_loss([num_positives, box_targets], box_outputs))
self.assertAllEqual(
legacy_fn.focal_loss(box_outputs, box_targets, alpha, gamma,
num_positives),
focal_loss_v2([num_positives, box_targets], box_outputs))
# TODO(tanmingxing): Re-enable this test after fixing this failing test.
# self.assertEqual(
# legacy_fn._box_iou_loss(box_outputs, box_targets, num_positives,
# 'ciou'),
# box_iou_loss([num_positives, box_targets], box_outputs))
iou_loss = box_iou_loss([num_positives, box_targets], box_outputs)
self.assertAlmostEqual(iou_loss.numpy(), 4.924635, places=5)
def _build_model(self, grad_checkpoint=False):
tf.random.set_seed(1111)
config = hparams_config.get_detection_config('efficientdet-d0')
config.heads = ['object_detection', 'segmentation']
config.batch_size = 1
config.num_examples_per_epoch = 1
config.model_dir = tempfile.mkdtemp()
config.steps_per_epoch = 1
config.mixed_precision = True
config.grad_checkpoint = grad_checkpoint
x = tf.ones((1, 512, 512, 3))
labels = {
'box_targets_%d' % i: tf.ones((1, 512 // 2**i, 512 // 2**i, 36))
for i in range(3, 8)
}
labels.update({
'cls_targets_%d' % i: tf.ones((1, 512 // 2**i, 512 // 2**i, 9),
dtype=tf.int32)
for i in range(3, 8)
})
labels.update({'image_masks': tf.ones((1, 128, 128, 1))})
labels.update({'mean_num_positives': tf.constant([10.0])})
params = config.as_dict()
params['num_shards'] = 1
params['steps_per_execution'] = 100
params['model_dir'] = tempfile.mkdtemp()
params['profile'] = False
config.override(params, allow_new_keys=True)
model = train_lib.EfficientDetNetTrain(config=config)
model.build((1, 512, 512, 3))
model.compile(
optimizer=train_lib.get_optimizer(params),
loss={
train_lib.BoxLoss.__name__:
train_lib.BoxLoss(params['delta'],
reduction=tf.keras.losses.Reduction.NONE),
train_lib.BoxIouLoss.__name__:
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),
train_lib.FocalLoss.__name__:
train_lib.FocalLoss(params['alpha'],
params['gamma'],
label_smoothing=params['label_smoothing'],
reduction=tf.keras.losses.Reduction.NONE),
tf.keras.losses.SparseCategoricalCrossentropy.__name__:
tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
})
return params, x, labels, model
def test_train(self):
tf.random.set_seed(1111)
config = hparams_config.get_detection_config('efficientdet-d0')
config.batch_size = 1
config.num_examples_per_epoch = 1
config.model_dir = tempfile.mkdtemp()
x = tf.ones((1, 512, 512, 3))
labels = {
'box_targets_%d' % i: tf.ones((1, 512 // 2**i, 512 // 2**i, 36))
for i in range(3, 8)
}
labels.update({
'cls_targets_%d' % i: tf.ones((1, 512 // 2**i, 512 // 2**i, 9),
dtype=tf.int32)
for i in range(3, 8)
})
labels.update({'mean_num_positives': tf.constant([10.0])})
params = config.as_dict()
params['num_shards'] = 1
model = train_lib.EfficientDetNetTrain(config=config)
model.build((1, 512, 512, 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'],
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)
})
# Test single-batch
outputs = model.train_on_batch(x, labels, return_dict=True)
self.assertAllClose(outputs, {'loss': 26278.2539}, rtol=.1, atol=100.)
outputs = model.test_on_batch(x, labels, return_dict=True)
self.assertAllClose(outputs, {'loss': 26061.1582}, rtol=.1, atol=100.)
# Test fit.
hist = model.fit(x,
labels,
steps_per_epoch=1,
epochs=1,
callbacks=train_lib.get_callbacks(params))
self.assertAllClose(hist.history, {'loss': [26061.1582]},
rtol=.1,
atol=100.)
def test_losses(self):
box_loss = train_lib.BoxLoss()
box_iou_loss = train_lib.BoxIouLoss('ciou')
alpha = 0.25
gamma = 1.5
focal_loss_v2 = train_lib.FocalLoss(
alpha, gamma, reduction=tf.keras.losses.Reduction.NONE)
box_outputs = tf.ones([8])
box_targets = tf.zeros([8])
num_positives = 4.0
self.assertEqual(
legacy_fn._box_loss(box_outputs, box_targets, num_positives),
box_loss([num_positives, box_targets], box_outputs))
self.assertEqual(
legacy_fn._box_iou_loss(box_outputs, box_targets, num_positives,
'ciou'),
box_iou_loss([num_positives, box_targets], box_outputs))
self.assertAllEqual(
legacy_fn.focal_loss(box_outputs, box_targets, alpha, gamma,
num_positives),
focal_loss_v2([num_positives, box_targets], box_outputs))
def test_train(self):
tf.random.set_seed(1111)
config = hparams_config.get_detection_config('efficientdet-d0')
config.batch_size = 1
config.num_examples_per_epoch = 1
config.model_dir = tempfile.mkdtemp()
x = tf.ones((1, 512, 512, 3))
labels = {
'box_targets_%d' % i: tf.ones((1, 512 // 2**i, 512 // 2**i, 36))
for i in range(3, 8)
}
labels.update({
'cls_targets_%d' % i: tf.ones((1, 512 // 2**i, 512 // 2**i, 9),
dtype=tf.int32) for i in range(3, 8)
})
labels.update({'mean_num_positives': tf.constant([10.0])})
params = config.as_dict()
params['num_shards'] = 1
model = train_lib.EfficientDetNetTrain(config=config)
model.build((1, 512, 512, 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'],
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)
})
# Test single-batch
outputs = model.train_on_batch(x, labels, return_dict=True)
expect_results = {'loss': 26278.25,
'det_loss': 26277.033203125,
'cls_loss': 5060.716796875,
'box_loss': 424.3263244628906,
'box_iou_loss': 0,
'gnorm': 5873.78759765625}
self.assertAllClose(outputs, expect_results, rtol=.1, atol=100.)
outputs = model.test_on_batch(x, labels, return_dict=True)
expect_results = {'loss': 26079.712890625,
'det_loss': 26078.49609375,
'cls_loss': 5063.3759765625,
'box_loss': 420.30242919921875,
'box_iou_loss': 0}
self.assertAllClose(outputs, expect_results, rtol=.1, atol=100.)
# Test fit.
hist = model.fit(
x,
labels,
steps_per_epoch=1,
epochs=1,
callbacks=train_lib.get_callbacks(params))
expect_results = {'loss': [26063.099609375],
'det_loss': [26061.8828125],
'cls_loss': [5058.1337890625],
'box_loss': [420.074951171875],
'box_iou_loss': [0],
'gnorm': [5107.46435546875]}
self.assertAllClose(
hist.history, expect_results, rtol=.1, atol=100.)
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 test_train(self):
tf.random.set_seed(1111)
config = hparams_config.get_detection_config('efficientdet-d0')
config.heads = ['object_detection', 'segmentation']
config.batch_size = 1
config.num_examples_per_epoch = 1
config.model_dir = tempfile.mkdtemp()
config.steps_per_epoch = 1
x = tf.ones((1, 512, 512, 3))
labels = {
'box_targets_%d' % i: tf.ones((1, 512 // 2**i, 512 // 2**i, 36))
for i in range(3, 8)
}
labels.update({
'cls_targets_%d' % i: tf.ones((1, 512 // 2**i, 512 // 2**i, 9),
dtype=tf.int32)
for i in range(3, 8)
})
labels.update({'image_masks': tf.ones((1, 128, 128, 1))})
labels.update({'mean_num_positives': tf.constant([10.0])})
params = config.as_dict()
params['num_shards'] = 1
model = train_lib.EfficientDetNetTrain(config=config)
model.build((1, 512, 512, 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),
'seg_loss':
tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
})
# Test single-batch
outputs = model.train_on_batch(x, labels, return_dict=True)
expect_results = {
'loss': [26278.3, 5061.9, 425.5, 1.217],
'det_loss': 26277.033203125,
'cls_loss': 5060.716796875,
'box_loss': 424.3263244628906,
'gnorm': 5873.78759765625,
'seg_loss': 1.2215478420257568,
}
self.assertAllClose(outputs, expect_results, rtol=.1, atol=100.)
outputs = model.test_on_batch(x, labels, return_dict=True)
expect_results = {
'loss': [26278.3, 5061.9, 425.5, 1.217],
'det_loss': 26078.49609375,
'cls_loss': 5063.3759765625,
'box_loss': 420.30242919921875,
'seg_loss': 1.2299377918243408,
}
self.assertAllClose(outputs, expect_results, rtol=.1, atol=100.)
# Test fit.
hist = model.fit(x,
labels,
steps_per_epoch=1,
epochs=1,
callbacks=train_lib.get_callbacks(params))
self.assertAllClose(hist.history['loss'],
[[26067, 5057.5, 421.4, 1.2]],
rtol=.1,
atol=10.)
self.assertAllClose(hist.history['det_loss'], [26061.],
rtol=.1,
atol=10.)
self.assertAllClose(hist.history['cls_loss'], [5058.],
rtol=.1,
atol=10.)
self.assertAllClose(hist.history['box_loss'], [420.],
rtol=.1,
atol=100.)
self.assertAllClose(hist.history['seg_loss'], [1.2299],
rtol=.1,
atol=100.)
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'))