[SparseZoo Integration] update sparsezoo usage for 1.1 refactor (#955)

* [SparseZoo v2 Bridge] Save IC training artifacts to ModelDirectory directory (#864)

* initial commit

* remove rubbish file

* Update train.py

* refactor export samples

* change name of folders

* working, time for refactoring

* ready for review

* Delete hehe.py

* Update helpers.py

* Apply suggestions from code review

Co-authored-by: bogunowicz@arrival.com <bogunowicz@arrival.com>
Co-authored-by: Benjamin Fineran <bfineran@users.noreply.github.com>

* first round of edits

* first sweep through all files

* checking if pytorch checks pass

* second sweep through all files

* third sweep through all files

* Update README.md

* Update README.md

* Update README.md

* fourth sweep through all files

* fifth sweep through all files

* Update integrations/keras/README.md

* rolling back some edits after bens comment

* remove notebook output

* add back onnx dataloading + batching functionality + tests

* pytorch logic, orgainization, completeness fixes

* use final checkpoint in eval mode IC training

Co-authored-by: bogunowicz@arrival.com <bogunowicz@arrival.com>
Co-authored-by: Benjamin Fineran <bfineran@users.noreply.github.com>
Co-authored-by: Benjamin <ben@neuralmagic.com>

Author: 4 people

integrations/keras/README.md

@@ -57,17 +57,18 @@ Complete lists are available online for all [models](https://sparsezoo.neuralmag
 
 Sample code for retrieving a model from the SparseZoo:
 ```python
-from sparsezoo import Zoo
+from sparsezoo import Model
 
-model = Zoo.load_model_from_stub("zoo:cv/classification/resnet_v1-50/keras/sparseml/imagenet/pruned-moderate")
+model = Model("zoo:cv/classification/resnet_v1-50/keras/sparseml/imagenet/pruned-moderate")
 print(model)
 ```
 
 Sample code for retrieving a recipe from the SparseZoo:
 ```python
-from sparsezoo import Zoo
+from sparsezoo import Model
 
-recipe = Zoo.load_recipe_from_stub("zoo:cv/classification/resnet_v1-50/keras/sparseml/imagenet/pruned-conservative/original")
+model = Model("zoo:cv/classification/resnet_v1-50/keras/sparseml/imagenet/pruned-conservative")
+recipe = model.recipes.default
 print(recipe)
 ```
 

integrations/keras/notebooks/classification.ipynb

@@ -55,13 +55,26 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "NameError",
+     "evalue": "name 'Zoo' is not defined",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
+      "Input \u001b[0;32mIn [2]\u001b[0m, in \u001b[0;36m<cell line: 37>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     34\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mRecipe file not found: \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(recipe_file_path))\n\u001b[1;32m     35\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m model_file_path, recipe_file_path\n\u001b[0;32m---> 37\u001b[0m model_file_path, recipe_file_path \u001b[38;5;241m=\u001b[39m \u001b[43mdownload_model_and_recipe\u001b[49m\u001b[43m(\u001b[49m\u001b[43mroot_dir\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     39\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mLoading model \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(model_file_path))\n\u001b[1;32m     40\u001b[0m model \u001b[38;5;241m=\u001b[39m keras\u001b[38;5;241m.\u001b[39mmodels\u001b[38;5;241m.\u001b[39mload_model(model_file_path)\n",
+      "Input \u001b[0;32mIn [2]\u001b[0m, in \u001b[0;36mdownload_model_and_recipe\u001b[0;34m(root_dir)\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     10\u001b[0m \u001b[38;5;124;03mDownload pretrained model and a pruning recipe\u001b[39;00m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     12\u001b[0m model_dir \u001b[38;5;241m=\u001b[39m os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mjoin(root_dir, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmnist\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m---> 13\u001b[0m zoo_model \u001b[38;5;241m=\u001b[39m \u001b[43mZoo\u001b[49m\u001b[38;5;241m.\u001b[39mload_model(\n\u001b[1;32m     14\u001b[0m         domain\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcv\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     15\u001b[0m         sub_domain\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mclassification\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     16\u001b[0m         architecture\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmnist\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     17\u001b[0m         sub_architecture\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     18\u001b[0m         framework\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mkeras\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     19\u001b[0m         repo\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msparseml\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     20\u001b[0m         dataset\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmnist\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     21\u001b[0m         training_scheme\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     22\u001b[0m         sparse_name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpruned\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     23\u001b[0m         sparse_category\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mconservative\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     24\u001b[0m         sparse_target\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     25\u001b[0m         override_parent_path\u001b[38;5;241m=\u001b[39mmodel_dir,\n\u001b[1;32m     26\u001b[0m     )\n\u001b[1;32m     27\u001b[0m zoo_model\u001b[38;5;241m.\u001b[39mdownload()\n\u001b[1;32m     29\u001b[0m model_file_path \u001b[38;5;241m=\u001b[39m zoo_model\u001b[38;5;241m.\u001b[39mframework_files[\u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39mdownloaded_path()\n",
+      "\u001b[0;31mNameError\u001b[0m: name 'Zoo' is not defined"
+     ]
+    }
+   ],
    "source": [
     "import os\n",
     "from sparseml.keras.utils import keras\n",
-    "from sparsezoo.models import Zoo\n",
+    "from sparsezoo.models import Model\n",
     "\n",
     "# Root directory for the notebook artifacts\n",
     "root_dir = \"./notebooks/keras\"\n",
@@ -71,26 +84,13 @@
     "    Download pretrained model and a pruning recipe\n",
     "    \"\"\"\n",
     "    model_dir = os.path.join(root_dir, \"mnist\")\n",
-    "    zoo_model = Zoo.load_model(\n",
-    "            domain=\"cv\",\n",
-    "            sub_domain=\"classification\",\n",
-    "            architecture=\"mnist\",\n",
-    "            sub_architecture=None,\n",
-    "            framework=\"keras\",\n",
-    "            repo=\"sparseml\",\n",
-    "            dataset=\"mnist\",\n",
-    "            training_scheme=None,\n",
-    "            sparse_name=\"pruned\",\n",
-    "            sparse_category=\"conservative\",\n",
-    "            sparse_target=None,\n",
-    "            override_parent_path=model_dir,\n",
-    "        )\n",
-    "    zoo_model.download()\n",
-    "\n",
-    "    model_file_path = zoo_model.framework_files[0].downloaded_path()\n",
+    "    zoo_model = Model(...)\n",
+    "\n",
+    "\n",
+    "    model_file_path = zoo_model.training.default.get_file(\"model.h5\").path\n",
     "    if not os.path.exists(model_file_path) or not model_file_path.endswith(\".h5\"):\n",
     "        raise RuntimeError(\"Model file not found: {}\".format(model_file_path))\n",
-    "    recipe_file_path = zoo_model.recipes[0].downloaded_path()\n",
+    "    recipe_file_path = zoo_model.recipes.default.path\n",
     "    if not os.path.exists(recipe_file_path):\n",
     "        raise RuntimeError(\"Recipe file not found: {}\".format(recipe_file_path))\n",
     "    return model_file_path, recipe_file_path\n",
@@ -424,9 +424,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python (keras_pruning)",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
-   "name": "keras_pruning"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -438,7 +438,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.9"
+   "version": "3.8.10"
   }
  },
  "nbformat": 4,

integrations/keras/prune_resnet20.py

@@ -37,7 +37,7 @@
 from sparseml.keras.utils.callbacks import LossesAndMetricsLoggingCallback
 from sparseml.keras.utils.exporter import ModelExporter
 from sparseml.keras.utils.logger import TensorBoardLogger
-from sparsezoo.models import Zoo
+from sparsezoo import Model
 
 
 # Root directory
@@ -64,33 +64,21 @@ def download_model_and_recipe(root_dir: str):
     """
     Download pretrained model and a pruning recipe
     """
-    model_dir = os.path.join(root_dir, "resnet20_v1")
 
-    # Load base model to prune
-    base_zoo_model = Zoo.load_model(
-        domain="cv",
-        sub_domain="classification",
-        architecture="resnet_v1",
-        sub_architecture=20,
-        framework="keras",
-        repo="sparseml",
-        dataset="cifar_10",
-        training_scheme=None,
-        sparse_name="base",
-        sparse_category="none",
-        sparse_target=None,
-        override_parent_path=model_dir,
-    )
-    base_zoo_model.download()
-    model_file_path = base_zoo_model.framework_files[0].downloaded_path()
-    if not os.path.exists(model_file_path) or not model_file_path.endswith(".h5"):
-        raise RuntimeError("Model file not found: {}".format(model_file_path))
-
-    # Simply use the recipe stub
+    # Use the recipe stub
     recipe_file_path = (
         "zoo:cv/classification/resnet_v1-20/keras/sparseml/cifar_10/pruned-conservative"
     )
 
+    # Load base model to prune
+    base_zoo_model = Model(recipe_file_path)
+    base_zoo_model.path = os.path.join(root_dir, "resnet20_v1")
+    checkpoint = base_zoo_model.training.default
+    model_file_path = checkpoint.get_file("model.h5").path
+    recipe_file_path = base_zoo_model.recipes.default.path
+    if not os.path.exists(model_file_path) or not model_file_path.endswith(".h5"):
+        raise RuntimeError("Model file not found: {}".format(model_file_path))
+
     return model_file_path, recipe_file_path
 
 

integrations/pytorch/README.md

@@ -60,17 +60,18 @@ Complete lists are available online for all [models](https://sparsezoo.neuralmag
 
 Sample code for retrieving a model from the SparseZoo:
 ```python
-from sparsezoo import Zoo
+from sparsezoo import Model
 
-model = Zoo.load_model_from_stub("zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned_quant-moderate")
+model = Model("zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned_quant-moderate")
 print(model)
 ```
 
 Sample code for retrieving a recipe from the SparseZoo:
 ```python
-from sparsezoo import Zoo
+from sparsezoo import Model
 
-recipe = Zoo.load_recipe_from_stub("zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned_quant-moderate/original")
+model = Model("zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned_quant-moderate")
+recipe = model.recipes.default
 print(recipe)
 ```
 

integrations/pytorch/notebooks/classification.ipynb

@@ -219,9 +219,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from sparsezoo import Zoo\n",
+    "from sparsezoo import Model, search_models\n",
     "\n",
-    "recipe = Zoo.search_recipes(\n",
+    "zoo_model = search_models(\n",
     "    domain=\"cv\",\n",
     "    sub_domain=\"classification\",\n",
     "    architecture=\"resnet_v1\",\n",
@@ -231,8 +231,7 @@
     "    dataset=\"imagenette\",\n",
     "    sparse_name=\"pruned\",\n",
     ")[0]  # unwrap search result\n",
-    "recipe.download()\n",
-    "recipe_path = recipe.downloaded_path()\n",
+    "recipe_path = zoo_model.recipes.default.path\n",
     "print(f\"Recipe downloaded to: {recipe_path}\")"
    ]
   },
@@ -364,4 +363,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

integrations/pytorch/notebooks/detection.ipynb

@@ -238,9 +238,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from sparsezoo import Zoo\n",
+    "from sparsezoo import Model, search_models\n",
     "\n",
-    "recipe = Zoo.search_recipes(\n",
+    "zoo_model = search_models(\n",
     "    domain=\"cv\",\n",
     "    sub_domain=\"detection\",\n",
     "    architecture=\"ssd\",\n",
@@ -250,8 +250,7 @@
     "    dataset=\"voc\",\n",
     "    sparse_name=\"pruned\",\n",
     ")[0]  # unwrap search result\n",
-    "recipe.download()\n",
-    "recipe_path = recipe.downloaded_path()\n",
+    "recipe_path = zoo_model.recipes.default.path\n",
     "print(f\"Recipe downloaded to: {recipe_path}\")"
    ]
   },
@@ -370,4 +369,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

integrations/pytorch/notebooks/sparse_quantized_transfer_learning.ipynb

@@ -76,7 +76,7 @@
    "source": [
     "from sparseml.pytorch.models import ModelRegistry\n",
     "from sparseml.pytorch.datasets import ImagenetteDataset, ImagenetteSize\n",
-    "from sparsezoo import Zoo\n",
+    "from sparsezoo import Model\n",
     "\n",
     "#######################################################\n",
     "# Define your model below\n",
@@ -85,7 +85,7 @@
     "# SparseZoo stub to pre-trained sparse-quantized ResNet-50 for imagenet dataset\n",
     "zoo_stub_path = (\n",
     "    \"zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned_quant-moderate\"\n",
-    "    \"?recipe_type=transfer_learn\"\n",
+    "    \"?recipe=transfer_learn\"\n",
     ")\n",
     "model = ModelRegistry.create(\n",
     "    key=\"resnet50\",\n",
@@ -236,9 +236,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from sparsezoo import Zoo\n",
+    "from sparsezoo import Model\n",
     "\n",
-    "recipe_path = Zoo.download_recipe_from_stub(zoo_stub_path)\n",
+    "zoo_model = Model(zoo_stub_path)\n",
+    "recipe_path = zoo_model.recipes.default.path\n",
     "print(f\"Recipe downloaded to: {recipe_path}\")"
    ]
   },
@@ -450,4 +451,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

integrations/pytorch/notebooks/torchvision.ipynb

@@ -270,9 +270,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from sparsezoo import Zoo\n",
+    "from sparsezoo import Model, search_models\n",
     "\n",
-    "recipe = Zoo.search_recipes(\n",
+    "zoo_model = search_models(\n",
     "    domain=\"cv\",\n",
     "    sub_domain=\"classification\",\n",
     "    architecture=\"resnet_v1\",\n",
@@ -282,8 +282,8 @@
     "    dataset=\"imagenette\",\n",
     "    sparse_name=\"pruned\",\n",
     ")[0]  # unwrap search result\n",
-    "recipe.download()\n",
-    "recipe_path = recipe.downloaded_path()\n",
+    "\n",
+    "recipe_path = zoo_model.recipes.default.path\n",
     "print(f\"Recipe downloaded to: {recipe_path}\")"
    ]
   },
@@ -396,4 +396,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

integrations/pytorch/utils.py

@@ -26,6 +26,9 @@
 from torch.utils.data import DataLoader
 
 from sparseml.pytorch.datasets import DatasetRegistry, ssd_collate_fn, yolo_collate_fn
+from sparseml.pytorch.image_classification.utils.helpers import (
+    download_framework_model_by_recipe_type,
+)
 from sparseml.pytorch.models import ModelRegistry
 from sparseml.pytorch.optim import ScheduledModifierManager, ScheduledOptimizer
 from sparseml.pytorch.sparsification import ConstantPruningModifier
@@ -45,7 +48,7 @@
     torch_distributed_zero_first,
 )
 from sparseml.utils import create_dirs
-from sparsezoo import Zoo
+from sparsezoo import Model
 
 
 @unique
@@ -237,9 +240,10 @@ def create_model(args: Any, num_classes: int) -> Module:
     with torch_distributed_zero_first(args.local_rank):  # only download once locally
         if args.checkpoint_path == "zoo":
             if args.recipe_path and args.recipe_path.startswith("zoo:"):
-                args.checkpoint_path = Zoo.download_recipe_base_framework_files(
-                    args.recipe_path, extensions=[".pth"]
-                )[0]
+                zoo_model = Model(args.recipe_path)
+                args.checkpoint_path = download_framework_model_by_recipe_type(
+                    zoo_model
+                )
             else:
                 raise ValueError(
                     "'zoo' provided as --checkpoint-path but a SparseZoo stub"

integrations/tensorflow_v1/README.md

@@ -57,17 +57,18 @@ Complete lists are available online for all [models](https://sparsezoo.neuralmag
 
 Sample code for retrieving a model from the SparseZoo:
 ```python
-from sparsezoo import Zoo
+from sparsezoo import Model
 
-model = Zoo.load_model_from_stub("zoo:cv/classification/resnet_v1-50/tensorflow_v1/sparseml/imagenette/pruned-moderate")
+model = Model("zoo:cv/classification/resnet_v1-50/tensorflow_v1/sparseml/imagenette/pruned-moderate")
 print(model)
 ```
 
 Sample code for retrieving a recipe from the SparseZoo:
 ```python
-from sparsezoo import Zoo
+from sparsezoo import Model
 
-recipe = Zoo.load_recipe_from_stub("zoo:cv/classification/resnet_v1-50/tensorflow_v1/sparseml/imagenette/pruned-moderate/original")
+model = Model("zoo:cv/classification/resnet_v1-50/tensorflow_v1/sparseml/imagenette/pruned-moderate")
+recipe = model.recipes.default
 print(recipe)
 ```