Revert "removed ScheduledOptimizer, moved logic to ScheduledModifierM…

README.md

@@ -102,29 +102,21 @@ For a more in-depth read, check out [SparseML documentation](https://docs.neural
 The PyTorch optimization libraries are located under the `sparseml.pytorch.optim` package.
 Inside are APIs designed to make model optimization as easy as possible by integrating seamlessly into PyTorch training pipelines.
 
-The integration is done using the `ScheduledModifierManager` class. It wraps your current optimizer's step function to apply SparseML `Modifier` optimizations at each step.
-The `ScheduledModifierManager` class can be created from a recipe file or `SparseZoo` optimized model stub.
-
-With this setup, the training process can then be modified as desired to optimize the model using SparseML recipes.
+The integration is done using the `ScheduledOptimizer` class. It is intended to wrap your current optimizer and its step function. The step function then calls into the `ScheduledModifierManager` class which can be created from a recipe file. With this setup, the training process can then be modified as desired to optimize the model.
 
 To enable all of this, the integration code you'll need to write is only a handful of lines:
 
 ```python
-from sparseml.pytorch.optim import ScheduledModifierManager
+from sparseml.pytorch.optim import ScheduledModifierManager, ScheduledOptimizer
 
-# setup
 model = None  # your model definition
 optimizer = None  # your optimizer definition
 num_train_batches = len(train_data) / batch_size  # your number of batches per training epoch
 
-# integration
 manager = ScheduledModifierManager.from_yaml("/PATH/TO/recipe.yaml")
-manager.initialize(model, optimizer, steps_per_epoch=num_train_batches)
+optimizer = ScheduledOptimizer(optimizer, model, manager, steps_per_epoch=num_train_batches)
 
 # PyTorch training code...
-
-# finalize cleans up any added parameters or hooks
-manager.finalize(model, optimizer)
 ```
 
 ### Keras Optimization

docs/source/quicktour.md

@@ -61,29 +61,21 @@ For a more in-depth read, check out [SparseML documentation](https://docs.neural
 The PyTorch optimization libraries are located under the `sparseml.pytorch.optim` package.
 Inside are APIs designed to make model optimization as easy as possible by integrating seamlessly into PyTorch training pipelines.
 
-The integration is done using the `ScheduledModifierManager` class. It wraps your current optimizer's step function to apply SparseML `Modifier` optimizations at each step.
-The `ScheduledModifierManager` class can be created from a recipe file or `SparseZoo` optimized model stub.
-
-With this setup, the training process can then be modified as desired to optimize the model using SparseML recipes.
+The integration is done using the `ScheduledOptimizer` class. It is intended to wrap your current optimizer and its step function. The step function then calls into the `ScheduledModifierManager` class which can be created from a recipe file. With this setup, the training process can then be modified as desired to optimize the model.
 
 To enable all of this, the integration code you'll need to write is only a handful of lines:
 
 ```python
-from sparseml.pytorch.optim import ScheduledModifierManager
+from sparseml.pytorch.optim import ScheduledModifierManager, ScheduledOptimizer
 
-# setup
 model = None  # your model definition
 optimizer = None  # your optimizer definition
 num_train_batches = len(train_data) / batch_size  # your number of batches per training epoch
 
-# integration
 manager = ScheduledModifierManager.from_yaml("/PATH/TO/recipe.yaml")
-manager.initialize(model, optimizer, steps_per_epoch=num_train_batches)
+optimizer = ScheduledOptimizer(optimizer, model, manager, steps_per_epoch=num_train_batches)
 
 # PyTorch training code...
-
-# finalize cleans up any added parameters or hooks
-manager.finalize(model, optimizer)
 ```
 
 ### Keras Optimization

examples/pytorch-torchvision/main.py

@@ -95,7 +95,7 @@
 from torchvision import models
 
 from sparseml.pytorch.datasets.classification import ImageFolderDataset
-from sparseml.pytorch.optim import ScheduledModifierManager
+from sparseml.pytorch.optim import ScheduledModifierManager, ScheduledOptimizer
 from sparseml.pytorch.utils import ModuleExporter, PythonLogger, load_model
 from sparseml.utils import create_dirs
 
@@ -408,9 +408,10 @@ def main(args):
     # add sparseml modifiers #
     ##########################
     manager = ScheduledModifierManager.from_yaml(args.recipe_path)
-    manager.initialize(
-        model,
+    optimizer = ScheduledOptimizer(
         optimizer,
+        model,
+        manager,
         steps_per_epoch=len(train_loader),
         loggers=[PythonLogger()],
     )
@@ -431,7 +432,6 @@ def main(args):
     ########################
     # export trained model #
     ########################
-    manager.finalize(model, optimizer)
     exporter = ModuleExporter(model, save_dir)
     sample_input = torch.randn(image_shape).unsqueeze(0)  # sample batch for ONNX export
     exporter.export_onnx(sample_input)

examples/pytorch-torchvision/pruning.ipynb

@@ -272,9 +272,18 @@
    "source": [
     "from sparsezoo import Zoo\n",
     "\n",
-    "recipe_path = Zoo.download_recipe_from_stub(\n",
-    "    \"zoo:cv/classification/resnet_v1-50/pytorch/torchvision/imagenette/pruned-conservative\"\n",
-    ")\n",
+    "recipe = Zoo.search_recipes(\n",
+    "    domain=\"cv\",\n",
+    "    sub_domain=\"classification\",\n",
+    "    architecture=\"resnet_v1\",\n",
+    "    sub_architecture=\"50\",\n",
+    "    framework=\"pytorch\",\n",
+    "    repo=\"torchvision\",\n",
+    "    dataset=\"imagenette\",\n",
+    "    optim_name=\"pruned\",\n",
+    ")[0]  # unwrap search result\n",
+    "recipe.download()\n",
+    "recipe_path = recipe.downloaded_path()\n",
     "print(f\"Recipe downloaded to: {recipe_path}\")"
    ]
   },
@@ -284,15 +293,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from sparseml.pytorch.optim import ScheduledModifierManager\n",
+    "from sparseml.pytorch.optim import (\n",
+    "    ScheduledModifierManager,\n",
+    "    ScheduledOptimizer,\n",
+    ")\n",
     "\n",
     "# create ScheduledModifierManager and Optimizer wrapper\n",
     "manager = ScheduledModifierManager.from_yaml(recipe_path)\n",
-    "manager.initialize(\n",
-    "    model,\n",
-    "    optimizer,\n",
-    "    steps_per_epoch=len(train_loader),\n",
-    "    loggers=[],\n",
+    "optimizer = ScheduledOptimizer(\n",
+    "    optimizer, model, manager, steps_per_epoch=len(train_loader), loggers=[],\n",
     ")\n",
     "\n",
     "train_model(\n",
@@ -303,9 +312,7 @@
     "    device,\n",
     "    num_epochs=manager.max_epochs,\n",
     "    is_inception=False,\n",
-    ")\n",
-    "\n",
-    "manager.finalize(model, optimizer)"
+    ")"
    ]
   },
   {

notebooks/pytorch_classification.ipynb

@@ -221,9 +221,18 @@
    "source": [
     "from sparsezoo import Zoo\n",
     "\n",
-    "recipe_path = Zoo.download_recipe_from_stub(\n",
-    "    \"zoo:cv/classification/resnet_v1-50/tensorflow_v1/sparseml/imagenette/pruned-moderate\"\n",
-    ")\n",
+    "recipe = Zoo.search_recipes(\n",
+    "    domain=\"cv\",\n",
+    "    sub_domain=\"classification\",\n",
+    "    architecture=\"resnet_v1\",\n",
+    "    sub_architecture=\"50\",\n",
+    "    framework=\"pytorch\",\n",
+    "    repo=\"sparseml\",\n",
+    "    dataset=\"imagenette\",\n",
+    "    optim_name=\"pruned\",\n",
+    ")[0]  # unwrap search result\n",
+    "recipe.download()\n",
+    "recipe_path = recipe.downloaded_path()\n",
     "print(f\"Recipe downloaded to: {recipe_path}\")"
    ]
   },
@@ -233,16 +242,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from sparseml.pytorch.optim import ScheduledModifierManager\n",
-    "\n",
+    "from sparseml.pytorch.optim import (\n",
+    "    ScheduledModifierManager,\n",
+    "    ScheduledOptimizer,\n",
+    ")\n",
     "\n",
     "# create ScheduledModifierManager and Optimizer wrapper\n",
     "manager = ScheduledModifierManager.from_yaml(recipe_path)\n",
-    "manager.initialize(\n",
-    "    model,\n",
-    "    optimizer,\n",
-    "    steps_per_epoch=len(train_loader),\n",
-    "    loggers=[],\n",
+    "optimizer = ScheduledOptimizer(\n",
+    "    optimizer, model, manager, steps_per_epoch=len(train_loader), loggers=[],\n",
     ")\n",
     "\n",
     "\n",
@@ -260,19 +268,19 @@
     "            epoch_name, train_loss, train_acc\n",
     "        )\n",
     "    )\n",
-    "\n",
+    "    \n",
     "    # run validation loop\n",
     "    print(\"Running Validation Epoch {}\".format(epoch_name))\n",
-    "    val_loss, val_acc = run_model_one_epoch(model, train_loader, criterion, device)\n",
+    "    val_loss, val_acc = run_model_one_epoch(\n",
+    "        model, train_loader, criterion, device\n",
+    "    )\n",
     "    print(\n",
     "        \"Validation Epoch: {}\\nVal Loss: {}\\nTop 1 Acc: {}\\n\".format(\n",
     "            epoch_name, val_loss, val_acc\n",
     "        )\n",
     "    )\n",
-    "\n",
-    "    epoch += 1\n",
-    "\n",
-    "manager.finalize()"
+    "    \n",
+    "    epoch += 1"
    ]
   },
   {

notebooks/pytorch_detection.ipynb

@@ -240,9 +240,18 @@
    "source": [
     "from sparsezoo import Zoo\n",
     "\n",
-    "recipe_path = Zoo.download_recipe_from_stub(\n",
-    "    \"zoo:cv/detection/ssd-resnet18_300/pytorch/sparseml/voc/pruned-moderate\"\n",
-    ")\n",
+    "recipe = Zoo.search_recipes(\n",
+    "    domain=\"cv\",\n",
+    "    sub_domain=\"detection\",\n",
+    "    architecture=\"ssd\",\n",
+    "    sub_architecture=\"resnet18_300\",\n",
+    "    framework=\"pytorch\",\n",
+    "    repo=\"sparseml\",\n",
+    "    dataset=\"voc\",\n",
+    "    optim_name=\"pruned\",\n",
+    ")[0]  # unwrap search result\n",
+    "recipe.download()\n",
+    "recipe_path = recipe.downloaded_path()\n",
     "print(f\"Recipe downloaded to: {recipe_path}\")"
    ]
   },
@@ -252,13 +261,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from sparseml.pytorch.optim import ScheduledModifierManager\n",
+    "from sparseml.pytorch.optim import (\n",
+    "    ScheduledModifierManager,\n",
+    "    ScheduledOptimizer,\n",
+    ")\n",
     "\n",
     "# create ScheduledModifierManager and Optimizer wrapper\n",
     "manager = ScheduledModifierManager.from_yaml(recipe_path)\n",
-    "manager.initialize(\n",
-    "    model,\n",
+    "optimizer = ScheduledOptimizer(\n",
     "    optimizer,\n",
+    "    model,\n",
+    "    manager,\n",
     "    steps_per_epoch=len(train_loader),\n",
     "    loggers=[],\n",
     ")\n",
@@ -278,9 +291,7 @@
     "    val_loss = run_model_one_epoch(model, train_loader, criterion, device)\n",
     "    print(\"Validation Epoch: {}\\nVal Loss: {}\\n\".format(epoch_name, val_loss))\n",
     "\n",
-    "    epoch += 1\n",
-    "    \n",
-    "manager.finalize()"
+    "    epoch += 1"
    ]
   },
   {

scripts/pytorch_vision.py

@@ -414,7 +414,7 @@
 import torch
 from torch.nn import Module
 from torch.nn import functional as torch_functional
-from torch.optim import SGD, Adam, Optimizer
+from torch.optim import SGD, Adam
 from torch.optim.optimizer import Optimizer
 from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
@@ -432,6 +432,7 @@
 from sparseml.pytorch.optim import (
     ConstantPruningModifier,
     ScheduledModifierManager,
+    ScheduledOptimizer,
     default_exponential_check_lrs,
     lr_loss_sensitivity,
     pruning_loss_sens_magnitude,
@@ -968,7 +969,7 @@ def _create_scheduled_optimizer(
     model: Module,
     train_loader: DataLoader,
     loggers: List[Any],
-) -> Tuple[int, Optimizer, ScheduledModifierManager]:
+) -> Tuple[int, ScheduledOptimizer, ScheduledModifierManager]:
     # optimizer setup
     if args.optim == "SGD":
         optim_const = SGD
@@ -1011,8 +1012,13 @@ def _create_scheduled_optimizer(
     manager = ScheduledModifierManager.from_yaml(
         file_path=args.recipe_path, add_modifiers=add_mods
     )
-    manager.initialize(model, optim, steps_per_epoch=len(train_loader), loggers=loggers)
-
+    optim = ScheduledOptimizer(
+        optim,
+        model,
+        manager,
+        steps_per_epoch=len(train_loader),
+        loggers=loggers,
+    )
     LOGGER.info("created manager: {}".format(manager))
     return epoch, optim, manager
 
@@ -1108,8 +1114,8 @@ def train(args, model, train_loader, val_loader, input_shape, save_dir, loggers)
         LOGGER.info("starting training from epoch {}".format(epoch))
 
         if epoch > 0:
-            LOGGER.info("adjusting Manager to restore point")
-            ScheduledModifierManager.adjust_optimizer_step(optim, epoch, 0)
+            LOGGER.info("adjusting ScheduledOptimizer to restore point")
+            optim.adjust_current_step(epoch, 0)
 
         best_loss = None
         val_res = None
@@ -1118,7 +1124,7 @@ def train(args, model, train_loader, val_loader, input_shape, save_dir, loggers)
             if args.debug_steps > 0:
                 # correct since all optimizer steps are not
                 # taken in the epochs for debug mode
-                ScheduledModifierManager.adjust_optimizer_step(optim, epoch, 0)
+                optim.adjust_current_step(epoch, 0)
 
             if args.rank != -1:  # sync DDP dataloaders
                 train_loader.sampler.set_epoch(epoch)
@@ -1167,8 +1173,6 @@ def train(args, model, train_loader, val_loader, input_shape, save_dir, loggers)
 
         # export the final model
         LOGGER.info("completed...")
-        # finalize manager
-        manager.finalize(model, optim)
         if args.is_main_process:
             _save_model_training(
                 model, optim, input_shape, "model", save_dir, epoch, val_res

src/sparseml/pytorch/optim/__init__.py

@@ -33,6 +33,7 @@
 from .modifier_pruning import *
 from .modifier_quantization import *
 from .modifier_regularizer import *
+from .optimizer import *
 from .sensitivity_as import *
 from .sensitivity_lr import *
 from .sensitivity_pruning import *