Refactor: torchvision save checkpoint code, to include the checkpoint (#1334)

Refactor: Model creation code, to try `ModelRegistry.create` even when `arch_key` is `None`; This tries to read the arch_key from the model checkpoint, if key is not present, the relevant function still raise the respective errors

Author: rahul-tuli

src/sparseml/pytorch/torchvision/train.py

@@ -368,7 +368,7 @@ def collate_fn(batch):
 
     _LOGGER.info("Creating model")
     local_rank = args.rank if args.distributed else None
-    model = _create_model(
+    model, arch_key = _create_model(
         arch_key=args.arch_key,
         local_rank=local_rank,
         pretrained=args.pretrained,
@@ -652,6 +652,7 @@ def log_metrics(tag: str, metrics: utils.MetricLogger, epoch: int, epoch_step: i
                 "state_dict": model_without_ddp.state_dict(),
                 "optimizer": optimizer.state_dict(),
                 "args": args,
+                "arch_key": arch_key,
             }
             if lr_scheduler:
                 checkpoint["lr_scheduler"] = lr_scheduler.state_dict()
@@ -703,7 +704,7 @@ def _create_model(
     device=None,
     num_classes=None,
 ):
-    if arch_key in ModelRegistry.available_keys():
+    if not arch_key or arch_key in ModelRegistry.available_keys():
         with torch_distributed_zero_first(local_rank):
             model = ModelRegistry.create(
                 key=arch_key,
@@ -712,6 +713,9 @@ def _create_model(
                 pretrained_dataset=pretrained_dataset,
                 num_classes=num_classes,
             )
+
+        if isinstance(model, tuple):
+            model, arch_key = model
     elif arch_key in torchvision.models.__dict__:
         # fall back to torchvision
         model = torchvision.models.__dict__[arch_key](
@@ -724,7 +728,7 @@ def _create_model(
             f"Unable to find {arch_key} in ModelRegistry or in torchvision.models"
         )
     model.to(device)
-    return model
+    return model, arch_key
 
 
 def _get_lr_scheduler(args, optimizer, checkpoint=None, manager=None):