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 _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_display_callback(self):
config = hparams_config.get_detection_config('efficientdet-d0')
config.batch_size = 1
config.num_examples_per_epoch = 1
config.model_dir = tempfile.mkdtemp()
sample_image = tf.ones([416, 416, 3])
display_callback = train_lib.DisplayCallback(sample_image)
model = train_lib.EfficientDetNetTrain(config=config)
model.build((1, 512, 512, 3))
display_callback.set_model(model)
display_callback.on_epoch_end(0, {})
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_display_callback(self):
config = hparams_config.get_detection_config('efficientdet-d0')
config.batch_size = 1
config.num_examples_per_epoch = 1
config.model_dir = tempfile.mkdtemp()
fake_image = tf.ones([512, 512, 3], dtype=tf.uint8)
fake_jpeg = tf.image.encode_jpeg(fake_image)
sample_image = 'ram://fake_image.jpg'
tf.io.write_file(sample_image, fake_jpeg)
display_callback = train_lib.DisplayCallback(sample_image, config.model_dir)
model = train_lib.EfficientDetNetTrain(config=config)
model.build((1, 512, 512, 3))
display_callback.set_model(model)
display_callback.on_epoch_end(0, {})
def test_predict(self):
x = np.random.random((1, 512, 512, 3)).astype(np.float32)
model = train_lib.EfficientDetNetTrain('efficientdet-d0')
cls_outputs, box_outputs = model(x)
self.assertLen(cls_outputs, 5)
self.assertLen(box_outputs, 5)
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.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,
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.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)
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 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'))
