def main(cfg):
######################################################################################
# Create Training Data
######################################################################################
datasets = build_dataset(cfg.data.train)
tf_datasets = [
build_dataloader(datasets,
cfg.batch_size_per_device,
cfg.workers_per_gpu,
num_gpus=hvd.size(),
dist=True)
]
######################################################################################
# Build Model
######################################################################################
model = build_detector(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
# Pass example through so tensor shapes are defined
model.CLASSES = datasets.CLASSES
_ = model(next(iter(tf_datasets[0][0])))
model.layers[0].layers[0].load_weights(cfg.weights_path, by_name=False)
######################################################################################
# Create Model Runner
######################################################################################
runner = sagemaker_runner.Runner(model,
batch_processor,
name=cfg.model_name,
optimizer=cfg.optimizer,
work_dir=cfg.work_dir,
logger=get_root_logger(cfg.log_level),
amp_enabled=cfg.fp16,
loss_weights=cfg.loss_weights)
runner.timestamp = int(time())
######################################################################################
# Setup Training Hooks
######################################################################################
runner.register_hook(
checkpoint.CheckpointHook(interval=cfg.checkpoint_interval,
out_dir=cfg.outputs_path,
s3_dir=None))
runner.register_hook(
CocoDistEvalmAPHook(cfg.data.val, interval=cfg.evaluation_interval))
runner.register_hook(iter_timer.IterTimerHook())
runner.register_hook(text.TextLoggerHook())
runner.register_hook(
visualizer.Visualizer(cfg.data.val, interval=100, top_k=10))
runner.register_hook(
tensorboard.TensorboardLoggerHook(log_dir=cfg.outputs_path,
interval=10,
image_interval=100,
s3_dir=None))
######################################################################################
# Run Model
######################################################################################
runner.run(tf_datasets, cfg.workflow, cfg.training_epochs)
def main(cfg):
decompress_data(cfg)
######################################################################################
# Create Training Data
######################################################################################
cfg.global_batch_size = cfg.batch_size_per_device * hvd.size()
cfg.steps_per_epoch = cfg.coco_images // cfg.global_batch_size
datasets = build_dataset(cfg.data.train)
tf_datasets = [
build_dataloader(datasets,
cfg.batch_size_per_device,
cfg.workers_per_gpu,
num_gpus=hvd.size(),
dist=True)
]
######################################################################################
# Build Model
######################################################################################
#update any hyperparams that we may have passed in via arguments
if cfg.ls > 0.0:
cfg.model['bbox_head']['label_smoothing'] = cfg.ls
if cfg.use_rcnn_bn:
cfg.model['bbox_head']['use_bn'] = cfg.use_rcnn_bn
if cfg.use_conv:
cfg.model['bbox_head']['use_conv'] = cfg.use_conv
cfg.schedule = args.schedule
model = build_detector(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
# Pass example through so tensor shapes are defined
_ = model(next(iter(tf_datasets[0][0])))
model.layers[0].layers[0].load_weights(cfg.weights_path, by_name=False)
######################################################################################
# Build optimizer and associate scheduler
######################################################################################
# base learning rate is set for global batch size of 8, with linear scaling for larger batches
base_learning_rate = cfg.base_learning_rate
scaled_learning_rate = base_learning_rate * cfg.global_batch_size / 8
steps_per_epoch = cfg.steps_per_epoch
if cfg.schedule == '1x':
scheduler = tf.keras.optimizers.schedules.PiecewiseConstantDecay(
[steps_per_epoch * 8, steps_per_epoch * 10], [
scaled_learning_rate, scaled_learning_rate * 0.1,
scaled_learning_rate * 0.01
])
elif cfg.schedule == 'cosine':
scheduler = tf.keras.experimental.CosineDecayRestarts(
initial_learning_rate=scaled_learning_rate,
first_decay_steps=12 * steps_per_epoch,
t_mul=1,
m_mul=1) #0-1-13
else:
raise NotImplementedError
warmup_init_lr = 1.0 / cfg.warmup_init_lr_scale * scaled_learning_rate
scheduler = WarmupScheduler(scheduler, warmup_init_lr, cfg.warmup_steps)
# FIXME: currently hardcoded to SGD
optimizer = tf.keras.optimizers.SGD(scheduler,
momentum=0.9,
nesterov=False)
if cfg.fp16:
optimizer = tf.train.experimental.enable_mixed_precision_graph_rewrite(
optimizer, loss_scale='dynamic')
######################################################################################
# Create Model Runner
######################################################################################
runner = sagemaker_runner.Runner(model,
batch_processor,
name=cfg.model_name,
optimizer=optimizer,
work_dir=cfg.work_dir,
logger=get_root_logger(cfg.log_level),
amp_enabled=cfg.fp16,
loss_weights=cfg.loss_weights)
runner.timestamp = int(time())
######################################################################################
# Setup Training Hooks
######################################################################################
runner.register_hook(
checkpoint.CheckpointHook(interval=cfg.checkpoint_interval,
out_dir=cfg.outputs_path,
s3_dir=cfg.s3_checkpoints,
h5=True))
runner.register_hook(
CocoDistEvalmAPHook(cfg.data.val, interval=cfg.evaluation_interval))
runner.register_hook(iter_timer.IterTimerHook())
runner.register_hook(text.TextLoggerHook())
runner.register_hook(
visualizer.Visualizer(cfg.data.val, interval=100, top_k=10))
runner.register_hook(
tensorboard.TensorboardLoggerHook(log_dir=cfg.outputs_path,
image_interval=100,
s3_dir=cfg.s3_tensorboard))
######################################################################################
# Run Model
######################################################################################
runner.run(tf_datasets, cfg.workflow, cfg.training_epochs)