def main(args):
# Use fixed random seed if requested.
if args.seed is not None:
print("USING FIX SEED: {}".format(args.seed))
ensure_reproducibility(seed=args.seed)
# Append timestamp to output directory.
args.output_dir = timestamp_path(args.output_dir)
args.output_dir.mkdir(parents=True)
# Development mode overrides.
args.log_interval = 1 if args.dev_mode else args.log_interval
args.max_epochs = 2 if args.dev_mode else args.max_epochs
args.backbone_epochs = 1 if args.dev_mode else args.backbone_epochs
image_mean, image_std = dataset_moments(args)
# Only use a subset of data in dev mode.
subset_indices = range(DEV_MODE_BATCHES) if args.dev_mode else None
train_dataloader, test_dataloader, num_classes = detection_dataloaders(
args.data_dir,
subset_indices=subset_indices,
no_augmentation=args.no_augmentation,
num_workers=args.num_workers,
)
with redirect_output():
coco_gt = convert_to_coco_api(test_dataloader.dataset)
model = initialize_detector(
num_classes,
args.dropout_rate,
trainable_backbone_layers=args.trainable_backbone_layers,
image_mean=image_mean,
image_std=image_std,
)
model = model.to(device=args.device)
optimizer, lr_scheduler = initialize_optimizer(model, args)
trainer = create_trainer(model, optimizer, args)
evaluator = create_evaluator(model, args)
wandb_mode = "offline" if args.dev_mode else None
wandb_logger = setup_wandb_logging(
trainer=trainer,
optimizers=optimizer,
evaluators=evaluator,
log_every_iters=args.log_interval,
# kwargs...
dir=args.output_dir,
mode=wandb_mode,
config=args,
config_exclude_keys=["dev_mode"],
group="detector",
)
if lr_scheduler is not None:
trainer.add_event_handler(
Events.EPOCH_COMPLETED,
lambda: lr_scheduler.step(),
)
if args.backbone_epochs is not None:
trainer.add_event_handler(
Events.EPOCH_COMPLETED(once=args.backbone_epochs),
freeze_backbone,
model,
)
@trainer.on(
Events.EPOCH_COMPLETED(every=args.checkpoint_interval)
| Events.COMPLETED)
def save_model(engine: Engine):
slug = f"{engine.state.epoch:02d}"
if engine.last_event_name is Events.COMPLETED:
slug = "final"
file = args.output_dir / f"ckpt-{slug}.pt"
torch.save(
{
"epoch": engine.state.epoch,
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"dropout_rate": args.dropout_rate,
"num_classes": num_classes,
}, file)
@trainer.on(Events.EPOCH_STARTED)
def log_epoch(engine: Engine):
wandb_logger.log({"epoch": engine.state.epoch}, commit=False)
@trainer.on(Events.ITERATION_COMPLETED(every=args.log_interval))
def log_training_step(engine: Engine):
engine.logger.info(engine.state.metrics)
@trainer.on(Events.EPOCH_COMPLETED)
def run_evaluator(engine: Engine):
evaluator.run(test_dataloader)
# COCO evaluation scores.
# -----------------------
@evaluator.on(Events.STARTED)
def prepare_coco_evaluator(engine: Engine):
iou_types = ["bbox", "segm"]
engine.state.coco_evaluator = CocoEvaluator(coco_gt, iou_types)
@evaluator.on(Events.ITERATION_COMPLETED)
def update_coco_evaluator(engine: Engine):
engine.state.coco_evaluator.update(engine.state.result)
@evaluator.on(Events.COMPLETED)
def log_coco_evaluator(engine: Engine):
with redirect_output():
engine.state.coco_evaluator.synchronize_between_processes()
engine.state.coco_evaluator.accumulate()
engine.state.coco_evaluator.summarize()
coco_scores = prepare_coco_scores(engine.state.coco_evaluator)
wandb_logger.log(coco_scores, step=trainer.state.iteration)
del engine.state.coco_evaluator
# Mask image visualization.
# -------------------------
@evaluator.on(Events.STARTED)
def prepare_mask_images(engine: Engine):
engine.state.result_images = []
@evaluator.on(Events.ITERATION_COMPLETED)
def visualize_masks(engine: Engine):
if len(engine.state.result_images) < MAX_MASK_IMAGES:
image = engine.state.batch[0][0] # Grab first image
result_image, *_ = predict(image, model, args)
engine.state.result_images.append(result_image)
@evaluator.on(Events.COMPLETED)
def log_masked_images(engine: Engine):
images = engine.state.result_images[:MAX_MASK_IMAGES]
wandb_logger.log({"images": [wandb.Image(img) for img in images]},
step=trainer.state.iteration)
del engine.state.result_images
# Start training procedure...
trainer.run(train_dataloader, max_epochs=args.max_epochs)
def test_setup_wandb_logging(dirname):
from unittest.mock import patch
with patch("ignite.contrib.engines.common.WandBLogger") as _:
setup_wandb_logging(MagicMock())
def get_logger(
config: Any,
trainer: Engine,
evaluator: Optional[Union[Engine, Dict[str, Engine]]] = None,
optimizers: Optional[Union[Optimizer, Dict[str, Optimizer]]] = None,
**kwargs: Any,
) -> Optional[BaseLogger]:
"""Get Ignite provided logger.
Parameters
----------
config
Config object for setting up loggers
`config` has to contain
- `filepath`: logging path to output file
- `logger_log_every_iters`: logging iteration interval for loggers
trainer
trainer engine
evaluator
evaluator engine
optimizers
optimizers to log optimizer parameters
kwargs
optional keyword arguments passed to the logger
Returns
-------
logger_handler
Ignite provided logger instance
"""
{% if logger_deps == 'clearml' %}
logger_handler = common.setup_clearml_logging(
trainer=trainer,
optimizers=optimizers,
evaluators=evaluator,
log_every_iters=config.logger_log_every_iters,
**kwargs,
)
{% elif logger_deps == 'mlflow' %}
logger_handler = common.setup_mlflow_logging(
trainer=trainer,
optimizers=optimizers,
evaluators=evaluator,
log_every_iters=config.logger_log_every_iters,
**kwargs,
)
{% elif logger_deps == 'neptune-client' %}
logger_handler = common.setup_neptune_logging(
trainer=trainer,
optimizers=optimizers,
evaluators=evaluator,
log_every_iters=config.logger_log_every_iters,
**kwargs,
)
{% elif logger_deps == 'polyaxon-client' %}
logger_handler = common.setup_plx_logging(
trainer=trainer,
optimizers=optimizers,
evaluators=evaluator,
log_every_iters=config.logger_log_every_iters,
**kwargs,
)
{% elif logger_deps == 'tensorboard' %}
logger_handler = common.setup_tb_logging(
output_path=config.output_dir,
trainer=trainer,
optimizers=optimizers,
evaluators=evaluator,
log_every_iters=config.logger_log_every_iters,
**kwargs,
)
{% elif logger_deps == 'visdom' %}
logger_handler = common.setup_visdom_logging(
trainer=trainer,
optimizers=optimizers,
evaluators=evaluator,
log_every_iters=config.logger_log_every_iters,
**kwargs,
)
{% elif logger_deps == 'wandb' %}
logger_handler = common.setup_wandb_logging(
trainer=trainer,
optimizers=optimizers,
evaluators=evaluator,
log_every_iters=config.logger_log_every_iters,
**kwargs,
)
{% else %}
logger_handler = None
{% endif %}
return logger_handler
