Updated cifar10 example (#1632)

* Updates for cifar10 example

* Updates for cifar10 example

* More updates

* Updated code

* Fixed code-formatting

Author: vfdev-5

examples/contrib/cifar10/main.py

@@ -6,13 +6,14 @@
 import torch.nn as nn
 import torch.optim as optim
 import utils
+from torch.cuda.amp import GradScaler, autocast
 
 import ignite
 import ignite.distributed as idist
 from ignite.contrib.engines import common
 from ignite.contrib.handlers import PiecewiseLinear
 from ignite.engine import Engine, Events, create_supervised_evaluator
-from ignite.handlers import Checkpoint, DiskSaver
+from ignite.handlers import Checkpoint, DiskSaver, global_step_from_engine
 from ignite.metrics import Accuracy, Loss
 from ignite.utils import manual_seed, setup_logger
 
@@ -76,8 +77,8 @@ def training(local_rank, config):
 
     # Let's now setup evaluator engine to perform model's validation and compute metrics
     metrics = {
-        "accuracy": Accuracy(),
-        "loss": Loss(criterion),
+        "Accuracy": Accuracy(),
+        "Loss": Loss(criterion),
     }
 
     # We define two evaluators as they wont have exactly similar roles:
@@ -102,15 +103,18 @@ def run_validation(engine):
         evaluators = {"training": train_evaluator, "test": evaluator}
         tb_logger = common.setup_tb_logging(output_path, trainer, optimizer, evaluators=evaluators)
 
-    # Store 3 best models by validation accuracy:
-    common.gen_save_best_models_by_val_score(
-        save_handler=get_save_handler(config),
-        evaluator=evaluator,
-        models={"model": model},
-        metric_name="accuracy",
-        n_saved=3,
-        trainer=trainer,
-        tag="test",
+    # Store 2 best models by validation accuracy starting from num_epochs / 2:
+    best_model_handler = Checkpoint(
+        {"model": model},
+        get_save_handler(config),
+        filename_prefix="best",
+        n_saved=2,
+        global_step_transform=global_step_from_engine(trainer),
+        score_name="test_accuracy",
+        score_function=Checkpoint.get_default_score_fn("accuracy"),
+    )
+    evaluator.add_event_handler(
+        Events.COMPLETED(lambda *_: trainer.state.epoch > config["num_epochs"] // 2), best_model_handler
     )
 
     # In order to check training resuming we can stop training on a given iteration
@@ -124,9 +128,8 @@ def _():
     try:
         trainer.run(train_loader, max_epochs=config["num_epochs"])
     except Exception as e:
-        import traceback
-
-        print(traceback.format_exc())
+        logger.exception("")
+        raise e
 
     if rank == 0:
         tb_logger.close()
@@ -145,13 +148,14 @@ def run(
     learning_rate=0.4,
     num_warmup_epochs=4,
     validate_every=3,
-    checkpoint_every=200,
+    checkpoint_every=1000,
     backend=None,
     resume_from=None,
     log_every_iters=15,
     nproc_per_node=None,
     stop_iteration=None,
     with_clearml=False,
+    with_amp=False,
     **spawn_kwargs,
 ):
     """Main entry to train an model on CIFAR10 dataset.
@@ -179,6 +183,7 @@ def run(
             It can be 0 to disable it. Default, 15.
         stop_iteration (int, optional): iteration to stop the training. Can be used to check resume from checkpoint.
         with_clearml (bool): if True, experiment ClearML logger is setup. Default, False.
+        with_amp (bool): if True, enables native automatic mixed precision. Default, False.
         **spawn_kwargs: Other kwargs to spawn run in child processes: master_addr, master_port, node_rank, nnodes
 
     """
@@ -245,13 +250,17 @@ def initialize(config):
 
 def log_metrics(logger, epoch, elapsed, tag, metrics):
     metrics_output = "\n".join([f"\t{k}: {v}" for k, v in metrics.items()])
-    logger.info(f"\nEpoch {epoch} - Evaluation time (seconds): {int(elapsed)} - {tag} metrics:\n {metrics_output}")
+    logger.info(f"\nEpoch {epoch} - Evaluation time (seconds): {elapsed:.2f} - {tag} metrics:\n {metrics_output}")
 
 
 def log_basic_info(logger, config):
     logger.info(f"Train {config['model']} on CIFAR10")
     logger.info(f"- PyTorch version: {torch.__version__}")
     logger.info(f"- Ignite version: {ignite.__version__}")
+    if torch.cuda.is_available():
+        logger.info(f"- GPU Device: {torch.cuda.get_device_name(idist.get_local_rank())}")
+        logger.info(f"- CUDA version: {torch.version.cuda}")
+        logger.info(f"- CUDNN version: {torch.backends.cudnn.version()}")
 
     logger.info("\n")
     logger.info("Configuration:")
@@ -279,6 +288,9 @@ def create_trainer(model, optimizer, criterion, lr_scheduler, train_sampler, con
     #    - RunningAverage` on `train_step` output
     #    - Two progress bars on epochs and optionally on iterations
 
+    with_amp = config["with_amp"]
+    scaler = GradScaler(enabled=with_amp)
+
     def train_step(engine, batch):
 
         x, y = batch[0], batch[1]
@@ -288,28 +300,21 @@ def train_step(engine, batch):
             y = y.to(device, non_blocking=True)
 
         model.train()
-        # Supervised part
-        y_pred = model(x)
-        loss = criterion(y_pred, y)
 
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
+        with autocast(enabled=with_amp):
+            y_pred = model(x)
+            loss = criterion(y_pred, y)
 
-        # This can be helpful for XLA to avoid performance slow down if fetch loss.item() every iteration
-        if config["log_every_iters"] > 0 and (engine.state.iteration - 1) % config["log_every_iters"] == 0:
-            batch_loss = loss.item()
-            engine.state.saved_batch_loss = batch_loss
-        else:
-            batch_loss = engine.state.saved_batch_loss
+        optimizer.zero_grad()
+        scaler.scale(loss).backward()
+        scaler.step(optimizer)
+        scaler.update()
 
         return {
-            "batch loss": batch_loss,
+            "batch loss": loss.item(),
         }
 
     trainer = Engine(train_step)
-    trainer.state.saved_batch_loss = -1.0
-    trainer.state_dict_user_keys.append("saved_batch_loss")
     trainer.logger = logger
 
     to_save = {"trainer": trainer, "model": model, "optimizer": optimizer, "lr_scheduler": lr_scheduler}

examples/contrib/cifar10_qat/main.py

@@ -6,13 +6,14 @@
 import torch.nn as nn
 import torch.optim as optim
 import utils
+from torch.cuda.amp import GradScaler, autocast
 
 import ignite
 import ignite.distributed as idist
 from ignite.contrib.engines import common
 from ignite.contrib.handlers import PiecewiseLinear
 from ignite.engine import Engine, Events, create_supervised_evaluator
-from ignite.handlers import Checkpoint, DiskSaver
+from ignite.handlers import Checkpoint, DiskSaver, global_step_from_engine
 from ignite.metrics import Accuracy, Loss
 from ignite.utils import manual_seed, setup_logger
 
@@ -31,16 +32,37 @@ def training(local_rank, config):
     if rank == 0:
         now = datetime.now().strftime("%Y%m%d-%H%M%S")
 
-        folder_name = "{}_backend-{}-{}_{}".format(config["model"], idist.backend(), idist.get_world_size(), now)
+        folder_name = f"{config['model']}_backend-{idist.backend()}-{idist.get_world_size()}_{now}"
         output_path = Path(output_path) / folder_name
         if not output_path.exists():
             output_path.mkdir(parents=True)
         config["output_path"] = output_path.as_posix()
-        logger.info("Output path: {}".format(config["output_path"]))
+        logger.info(f"Output path: {config['output_path']}")
 
         if "cuda" in device.type:
             config["cuda device name"] = torch.cuda.get_device_name(local_rank)
 
+        if config["with_clearml"]:
+            try:
+                from clearml import Task
+            except ImportError:
+                # Backwards-compatibility for legacy Trains SDK
+                from trains import Task
+
+            task = Task.init("CIFAR10-Training", task_name=output_path.stem)
+            task.connect_configuration(config)
+            # Log hyper parameters
+            hyper_params = [
+                "model",
+                "batch_size",
+                "momentum",
+                "weight_decay",
+                "num_epochs",
+                "learning_rate",
+                "num_warmup_epochs",
+            ]
+            task.connect({k: config[k] for k in hyper_params})
+
     # Setup dataflow, model, optimizer, criterion
     train_loader, test_loader = get_dataflow(config)
 
@@ -78,15 +100,18 @@ def run_validation(engine):
         evaluators = {"training": train_evaluator, "test": evaluator}
         tb_logger = common.setup_tb_logging(output_path, trainer, optimizer, evaluators=evaluators)
 
-    # Store 3 best models by validation accuracy:
-    common.save_best_model_by_val_score(
-        output_path=config["output_path"],
-        evaluator=evaluator,
-        model=model,
-        metric_name="Accuracy",
-        n_saved=1,
-        trainer=trainer,
-        tag="test",
+    # Store 2 best models by validation accuracy starting from num_epochs / 2:
+    best_model_handler = Checkpoint(
+        {"model": model},
+        get_save_handler(config),
+        filename_prefix="best",
+        n_saved=2,
+        global_step_transform=global_step_from_engine(trainer),
+        score_name="test_accuracy",
+        score_function=Checkpoint.get_default_score_fn("accuracy"),
+    )
+    evaluator.add_event_handler(
+        Events.COMPLETED(lambda *_: trainer.state.epoch > config["num_epochs"] // 2), best_model_handler
     )
 
     trainer.run(train_loader, max_epochs=config["num_epochs"])
@@ -108,11 +133,13 @@ def run(
     learning_rate=0.4,
     num_warmup_epochs=4,
     validate_every=3,
-    checkpoint_every=200,
+    checkpoint_every=1000,
     backend=None,
     resume_from=None,
     log_every_iters=15,
     nproc_per_node=None,
+    with_clearml=False,
+    with_amp=False,
     **spawn_kwargs,
 ):
     """Main entry to train an model on CIFAR10 dataset.
@@ -138,6 +165,8 @@ def run(
         resume_from (str, optional): path to checkpoint to use to resume the training from. Default, None.
         log_every_iters (int): argument to log batch loss every ``log_every_iters`` iterations.
             It can be 0 to disable it. Default, 15.
+        with_clearml (bool): if True, experiment ClearML logger is setup. Default, False.
+        with_amp (bool): if True, enables native automatic mixed precision. Default, False.
         **spawn_kwargs: Other kwargs to spawn run in child processes: master_addr, master_port, node_rank, nnodes
 
     """
@@ -149,10 +178,8 @@ def run(
     spawn_kwargs["nproc_per_node"] = nproc_per_node
 
     with idist.Parallel(backend=backend, **spawn_kwargs) as parallel:
-        try:
-            parallel.run(training, config)
-        except Exception as e:
-            raise e
+
+        parallel.run(training, config)
 
 
 def get_dataflow(config):
@@ -167,7 +194,7 @@ def get_dataflow(config):
         # Ensure that only rank 0 download the dataset
         idist.barrier()
 
-    # Setup data loader also adapted to distributed config
+    # Setup data loader also adapted to distributed config: nccl, gloo, xla-tpu
     train_loader = idist.auto_dataloader(
         train_dataset, batch_size=config["batch_size"], num_workers=config["num_workers"], shuffle=True, drop_last=True,
     )
@@ -180,6 +207,7 @@ def get_dataflow(config):
 
 def initialize(config):
     model = utils.get_model(config["model"])
+    # Adapt model for distributed settings if configured
     model = idist.auto_model(model, find_unused_parameters=True)
 
     optimizer = optim.SGD(
@@ -205,24 +233,28 @@ def initialize(config):
 
 def log_metrics(logger, epoch, elapsed, tag, metrics):
     metrics_output = "\n".join([f"\t{k}: {v}" for k, v in metrics.items()])
-    logger.info(f"\nEpoch {epoch} - Time taken (seconds) : {elapsed:.02f} - {tag} metrics:\n {metrics_output}")
+    logger.info(f"\nEpoch {epoch} - Evaluation time (seconds): {elapsed:.2f} - {tag} metrics:\n {metrics_output}")
 
 
 def log_basic_info(logger, config):
-    logger.info("Quantization Aware Training {} on CIFAR10".format(config["model"]))
-    logger.info("- PyTorch version: {}".format(torch.__version__))
-    logger.info("- Ignite version: {}".format(ignite.__version__))
+    logger.info(f"Quantization Aware Training {config['model']} on CIFAR10")
+    logger.info(f"- PyTorch version: {torch.__version__}")
+    logger.info(f"- Ignite version: {ignite.__version__}")
+    if torch.cuda.is_available():
+        logger.info(f"- GPU Device: {torch.cuda.get_device_name(idist.get_local_rank())}")
+        logger.info(f"- CUDA version: {torch.version.cuda}")
+        logger.info(f"- CUDNN version: {torch.backends.cudnn.version()}")
 
     logger.info("\n")
     logger.info("Configuration:")
     for key, value in config.items():
-        logger.info("\t{}: {}".format(key, value))
+        logger.info(f"\t{key}: {value}")
     logger.info("\n")
 
     if idist.get_world_size() > 1:
         logger.info("\nDistributed setting:")
-        logger.info("\tbackend: {}".format(idist.backend()))
-        logger.info("\tworld size: {}".format(idist.get_world_size()))
+        logger.info(f"\tbackend: {idist.backend()}")
+        logger.info(f"\tworld size: {idist.get_world_size()}")
         logger.info("\n")
 
 
@@ -239,6 +271,9 @@ def create_trainer(model, optimizer, criterion, lr_scheduler, train_sampler, con
     #    - RunningAverage` on `train_step` output
     #    - Two progress bars on epochs and optionally on iterations
 
+    with_amp = config["with_amp"]
+    scaler = GradScaler(enabled=with_amp)
+
     def train_step(engine, batch):
 
         x, y = batch[0], batch[1]
@@ -248,12 +283,15 @@ def train_step(engine, batch):
             y = y.to(device, non_blocking=True)
 
         model.train()
-        y_pred = model(x)
-        loss = criterion(y_pred, y)
+
+        with autocast(enabled=with_amp):
+            y_pred = model(x)
+            loss = criterion(y_pred, y)
 
         optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
+        scaler.scale(loss).backward()
+        scaler.step(optimizer)
+        scaler.update()
 
         return {
             "batch loss": loss.item(),
@@ -272,7 +310,7 @@ def train_step(engine, batch):
         train_sampler=train_sampler,
         to_save=to_save,
         save_every_iters=config["checkpoint_every"],
-        output_path=config["output_path"],
+        save_handler=get_save_handler(config),
         lr_scheduler=lr_scheduler,
         output_names=metric_names if config["log_every_iters"] > 0 else None,
         with_pbars=False,
@@ -282,13 +320,22 @@ def train_step(engine, batch):
     resume_from = config["resume_from"]
     if resume_from is not None:
         checkpoint_fp = Path(resume_from)
-        assert checkpoint_fp.exists(), "Checkpoint '{}' is not found".format(checkpoint_fp.as_posix())
-        logger.info("Resume from a checkpoint: {}".format(checkpoint_fp.as_posix()))
+        assert checkpoint_fp.exists(), f"Checkpoint '{checkpoint_fp.as_posix()}' is not found"
+        logger.info(f"Resume from a checkpoint: {checkpoint_fp.as_posix()}")
         checkpoint = torch.load(checkpoint_fp.as_posix(), map_location="cpu")
         Checkpoint.load_objects(to_load=to_save, checkpoint=checkpoint)
 
     return trainer
 
 
+def get_save_handler(config):
+    if config["with_clearml"]:
+        from ignite.contrib.handlers.clearml_logger import ClearMLSaver
+
+        return ClearMLSaver(dirname=config["output_path"])
+
+    return DiskSaver(config["output_path"], require_empty=False)
+
+
 if __name__ == "__main__":
     fire.Fire({"run": run})