Testing fixes - notebooks, registry, model architectures (#25)

notebooks/pytorch_detection.ipynb

@@ -81,6 +81,7 @@
     "    pretrained=True,\n",
     "    pretrained_dataset=\"voc\",\n",
     "    pretrained_backbone=False,  # no need to download initial weights\n",
+    "    num_classes=21,\n",
     ")\n",
     "model_name = model.__class__.__name__\n",
     "input_shape = ModelRegistry.input_shape(\"ssd300_resnet18\")\n",

notebooks/tensorflow_v1_classification.ipynb

@@ -178,7 +178,7 @@
     "        sub_domain=\"classification\",\n",
     "        architecture=\"resnet_v1\",\n",
     "        sub_architecture=\"50\",\n",
-    "        framework=\"pytorch\",\n",
+    "        framework=\"tensorflow_v1\",\n",
     "        repo=\"sparseml\",\n",
     "        dataset=\"imagenette\",\n",
     "        optim_name=\"pruned\",\n",
@@ -399,4 +399,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

src/sparseml/keras/utils/__init__.py

@@ -2,7 +2,7 @@
 Generic code used as utilities and helpers for Keras
 """
 
+from .callbacks import *
 from .exporter import *
 from .logger import *
-from .callbacks import *
 from .model import *

src/sparseml/keras/utils/callbacks.py

@@ -3,10 +3,11 @@
 """
 
 from typing import List, Union
+
 import tensorflow
-from tensorflow import keras
-from tensorflow import Tensor
-from sparseml.keras.utils import KerasLogger, LoggingMode
+from tensorflow import Tensor, keras
+
+from sparseml.keras.utils.logger import KerasLogger, LoggingMode
 
 
 __all__ = [
@@ -89,7 +90,7 @@ def on_test_batch_begin(self, batch, logs=None):
         Called at the begin of a batch in evaluation
 
         :param batch: batch index in current epoch
-        :param logs: dictionary of logs (see Keras Callback doc)        
+        :param logs: dictionary of logs (see Keras Callback doc)
         """
         super().on_test_batch_begin(batch, logs)
         self._set_logging_mode(LoggingMode.TEST)

src/sparseml/onnx/utils/__init__.py

@@ -4,8 +4,8 @@
 
 from .data import *
 from .graph_editor import *
+from .graph_optimizer import *
 from .helpers import *
 from .loss import *
 from .model import *
 from .sparse_tensor import *
-from .graph_optimizer import *

src/sparseml/pytorch/models/__init__.py

@@ -2,7 +2,6 @@
 Code for creating and loading models in PyTorch
 """
 
-from .registry import *
 from .classification import *
 from .detection import *
 from .external import *

src/sparseml/pytorch/models/classification/inception_v3.py

@@ -455,7 +455,7 @@ def forward(self, x_tens: Tensor) -> Tuple[Tensor, ...]:
     input_shape=(3, 299, 299),
     domain="cv",
     sub_domain="classification",
-    architecture="inception-v3",
+    architecture="inception_v3",
     sub_architecture="none",
     default_dataset="imagenet",
     default_desc="base",

src/sparseml/pytorch/models/classification/mobilenet.py

@@ -271,7 +271,7 @@ def _mobilenet_base_section_settings():
     input_shape=(3, 224, 224),
     domain="cv",
     sub_domain="classification",
-    architecture="mobilenet-v1",
+    architecture="mobilenet_v1",
     sub_architecture="1.0",
     default_dataset="imagenet",
     default_desc="base",
@@ -301,7 +301,7 @@ def mobilenet(num_classes: int = 1000, class_type: str = "single") -> MobileNet:
     input_shape=(3, 224, 224),
     domain="cv",
     sub_domain="classification",
-    architecture="mobilenet-v1",
+    architecture="mobilenet_v1",
     sub_architecture="han",
     default_dataset="imagenet",
     default_desc="base",

src/sparseml/pytorch/models/classification/mobilenet_v2.py

@@ -393,7 +393,7 @@ def mobilenet_v2_width(
     input_shape=(3, 224, 224),
     domain="cv",
     sub_domain="classification",
-    architecture="mobilenet-v2",
+    architecture="mobilenet_v2",
     sub_architecture="1.0",
     default_dataset="imagenet",
     default_desc="base",

src/sparseml/pytorch/models/classification/resnet.py

@@ -522,7 +522,7 @@ def create_section(settings: ResNetSectionSettings) -> Sequential:
     input_shape=(3, 224, 224),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v1",
+    architecture="resnet_v1",
     sub_architecture="18",
     default_dataset="imagenet",
     default_desc="base",
@@ -563,7 +563,7 @@ def resnet18(num_classes: int = 1000, class_type: str = "single") -> ResNet:
     input_shape=(3, 224, 224),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v2",
+    architecture="resnet_v2",
     sub_architecture="18",
     default_dataset="imagenet",
     default_desc="base",
@@ -604,7 +604,7 @@ def resnetv2_18(num_classes: int = 1000, class_type: str = "single") -> ResNet:
     input_shape=(3, 224, 224),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v1",
+    architecture="resnet_v1",
     sub_architecture="34",
     default_dataset="imagenet",
     default_desc="base",
@@ -645,7 +645,7 @@ def resnet34(num_classes: int = 1000, class_type: str = "single") -> ResNet:
     input_shape=(3, 224, 224),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v2",
+    architecture="resnet_v2",
     sub_architecture="34",
     default_dataset="imagenet",
     default_desc="base",
@@ -686,7 +686,7 @@ def resnetv2_34(num_classes: int = 1000, class_type: str = "single") -> ResNet:
     input_shape=(3, 224, 224),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v1",
+    architecture="resnet_v1",
     sub_architecture="50",
     default_dataset="imagenet",
     default_desc="base",
@@ -743,7 +743,7 @@ def resnet50(num_classes: int = 1000, class_type: str = "single") -> ResNet:
     input_shape=(3, 224, 224),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v2",
+    architecture="resnet_v2",
     sub_architecture="50",
     default_dataset="imagenet",
     default_desc="base",
@@ -810,7 +810,7 @@ def resnetv2_50(num_classes: int = 1000, class_type: str = "single") -> ResNet:
     input_shape=(3, 224, 224),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v1",
+    architecture="resnet_v1",
     sub_architecture="50-2xwidth",
     default_dataset="imagenet",
     default_desc="base",
@@ -928,7 +928,7 @@ def resnext50(num_classes: int = 1000, class_type: str = "single") -> ResNet:
     input_shape=(3, 224, 224),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v1",
+    architecture="resnet_v1",
     sub_architecture="101",
     default_dataset="imagenet",
     default_desc="base",
@@ -985,7 +985,7 @@ def resnet101(num_classes: int = 1000, class_type: str = "single") -> ResNet:
     input_shape=(3, 224, 224),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v2",
+    architecture="resnet_v2",
     sub_architecture="101",
     default_dataset="imagenet",
     default_desc="base",
@@ -1052,7 +1052,7 @@ def resnetv2_101(num_classes: int = 1000, class_type: str = "single") -> ResNet:
     input_shape=(3, 224, 224),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v1",
+    architecture="resnet_v1",
     sub_architecture="101-2xwidth",
     default_dataset="imagenet",
     default_desc="base",
@@ -1170,7 +1170,7 @@ def resnext101(num_classes: int = 1000, class_type: str = "single") -> ResNet:
     input_shape=(3, 224, 224),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v1",
+    architecture="resnet_v1",
     sub_architecture="152",
     default_dataset="imagenet",
     default_desc="base",
@@ -1227,7 +1227,7 @@ def resnet152(num_classes: int = 1000, class_type: str = "single") -> ResNet:
     input_shape=(3, 224, 224),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v2",
+    architecture="resnet_v2",
     sub_architecture="152",
     default_dataset="imagenet",
     default_desc="base",

src/sparseml/pytorch/models/detection/ssd_mobilenet.py

@@ -76,8 +76,8 @@ def get_feature_extractor(self) -> nn.Module:
     input_shape=(3, 300, 300),
     domain="cv",
     sub_domain="detection",
-    architecture="ssd-lite",
-    sub_architecture="mobilenet-v2",
+    architecture="ssd_lite",
+    sub_architecture="mobilenet_v2",
     default_dataset="coco",
     default_desc="base",
 )

src/sparseml/pytorch/models/detection/yolo_v3.py

@@ -262,8 +262,8 @@ def forward(self, inp: Tensor):
     input_shape=(3, 640, 640),
     domain="cv",
     sub_domain="detection",
-    architecture="yolo",
-    sub_architecture="v3",
+    architecture="yolo_v3",
+    sub_architecture="none",
     default_dataset="coco",
     default_desc="base",
 )

src/sparseml/pytorch/models/external/torchvision.py

@@ -69,7 +69,7 @@ def wrapper(
             if provided will override the pretrained param
         :param pretrained: True to load the default pretrained weights,
             a string to load a specific pretrained weight
-            (ex: base, optim, optim-perf),
+            (ex: base, pruned-moderate),
             or False to not load any pretrained weights
         :param pretrained_dataset: The dataset to load pretrained weights for
             (ex: imagenet, mnist, etc).
@@ -96,11 +96,13 @@ def wrapper(
                 key, pretrained, pretrained_dataset
             )
             try:
-                paths = zoo_model.download_framework_files()
+                paths = zoo_model.download_framework_files(extensions=[".pth"])
                 load_model(paths[0], model, load_strict, ignore)
             except Exception as ex:
                 # try one more time with overwrite on in case file was corrupted
-                paths = zoo_model.download_framework_files(overwrite=True)
+                paths = zoo_model.download_framework_files(
+                    overwrite=True, extensions=[".pth"]
+                )
                 load_model(paths[0], model, load_strict, ignore)
 
         return model

src/sparseml/pytorch/models/registry.py

@@ -69,7 +69,7 @@ def create(
 
         :param key: the model key (name) to create
         :param pretrained: True to load pretrained weights; to load a specific version
-            give a string with the name of the version (optim, optim-perf), default None
+            give a string with the name of the version (pruned-moderate, base), default None
         :param pretrained_path: A model file path to load into the created model
         :param pretrained_dataset: The dataset to load for the model
         :param load_strict: True to make sure all states are found in and
@@ -332,11 +332,13 @@ def wrapper(
                     key, pretrained, pretrained_dataset
                 )
                 try:
-                    paths = zoo_model.download_framework_files()
+                    paths = zoo_model.download_framework_files(extensions=[".pth"])
                     load_model(paths[0], model, load_strict, ignore)
                 except Exception as ex:
                     # try one more time with overwrite on in case file was corrupted
-                    paths = zoo_model.download_framework_files(overwrite=True)
+                    paths = zoo_model.download_framework_files(
+                        overwrite=True, extensions=[".pth"]
+                    )
                     load_model(paths[0], model, load_strict, ignore)
 
             return model

src/sparseml/tensorflow_v1/models/classification/mobilenet.py

@@ -264,7 +264,7 @@ def mobilenet_const(
     input_shape=(224, 224, 3),
     domain="cv",
     sub_domain="classification",
-    architecture="mobilenet-v1",
+    architecture="mobilenet_v1",
     sub_architecture="1.0",
     default_dataset="imagenet",
     default_desc="base",

src/sparseml/tensorflow_v1/models/classification/mobilenet_v2.py

@@ -490,7 +490,7 @@ def mobilenet_v2_width(
     input_shape=(224, 224, 3),
     domain="cv",
     sub_domain="classification",
-    architecture="mobilenet-v2",
+    architecture="mobilenet_v2",
     sub_architecture="1.0",
     default_dataset="imagenet",
     default_desc="base",

src/sparseml/tensorflow_v1/models/classification/resnet.py

@@ -431,7 +431,7 @@ def resnet_const(
     input_shape=(224, 224, 3),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v1",
+    architecture="resnet_v1",
     sub_architecture="18",
     default_dataset="imagenet",
     default_desc="base",
@@ -491,7 +491,7 @@ def resnet18(
     input_shape=(32, 32, 3),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v1",
+    architecture="resnet_v1",
     sub_architecture="20",
     default_dataset="Cifar10",
     default_desc="base",
@@ -537,7 +537,7 @@ def resnet20(
     input_shape=(224, 224, 3),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v1",
+    architecture="resnet_v1",
     sub_architecture="34",
     default_dataset="imagenet",
     default_desc="base",
@@ -597,7 +597,7 @@ def resnet34(
     input_shape=(224, 224, 3),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v1",
+    architecture="resnet_v1",
     sub_architecture="50",
     default_dataset="imagenet",
     default_desc="base",
@@ -677,7 +677,7 @@ def resnet50(
     input_shape=(224, 224, 3),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v1",
+    architecture="resnet_v1",
     sub_architecture="101",
     default_dataset="imagenet",
     default_desc="base",
@@ -757,7 +757,7 @@ def resnet101(
     input_shape=(224, 224, 3),
     domain="cv",
     sub_domain="classification",
-    architecture="resnet-v1",
+    architecture="resnet_v1",
     sub_architecture="152",
     default_dataset="imagenet",
     default_desc="base",

src/sparseml/tensorflow_v1/models/registry.py

@@ -120,7 +120,7 @@ def create_zoo_model(
 
         :param key: the model key (name) to retrieve
         :param pretrained: True to load pretrained weights; to load a specific version
-            give a string with the name of the version (optim, optim-perf), default True
+            give a string with the name of the version (pruned-moderate, base), default True
         :param pretrained_dataset: The dataset to load for the model
         :return: the sparsezoo Model reference for the given model
         """