Merge pull request #159 from nebuly-ai/add_save_function

Speedster: Add save function and edit notebooks

apps/accelerate/speedster/setup.py

@@ -3,7 +3,7 @@
 
 
 REQUIREMENTS = [
-    "nebullvm>=0.7.3",
+    "nebullvm>=0.8.0",
     "tabulate>=0.8.0",
 ]
 
@@ -12,7 +12,7 @@
 
 setup(
     name="speedster",
-    version="0.1.3",
+    version="0.2.0",
     packages=find_packages(),
     install_requires=REQUIREMENTS,
     long_description=long_description,

apps/accelerate/speedster/speedster/__init__.py

@@ -1 +1,5 @@
 from speedster.api.functions import optimize_model  # noqa: F401
+from nebullvm.operations.inference_learners.utils import (  # noqa: F401
+    load_model,
+    save_model,
+)

apps/accelerate/speedster/speedster/api/tests/test_huggingface.py

@@ -4,12 +4,11 @@
 from nebullvm.operations.inference_learners.huggingface import (
     HuggingFaceInferenceLearner,
 )
-from nebullvm.operations.optimizations.utils import load_model
 from nebullvm.optional_modules.tensorflow import tensorflow as tf
 from nebullvm.optional_modules.torch import torch
 from transformers import AlbertModel, TFAlbertModel, AlbertTokenizer
 
-from speedster import optimize_model
+from speedster import optimize_model, load_model
 
 
 def test_torch_huggingface_ort_input_text():
@@ -46,7 +45,7 @@ def test_torch_huggingface_ort_input_text():
         ),
     )
 
-    #save and load
+    # save and load
     with TemporaryDirectory() as tmp_dir:
         optimized_model.save(tmp_dir)
         loaded_model = load_model(tmp_dir)

apps/accelerate/speedster/speedster/api/tests/test_onnx.py

@@ -18,10 +18,9 @@
     NumpyApacheTVMInferenceLearner,
 )
 from nebullvm.operations.optimizations.compilers.utils import tvm_is_available
-from nebullvm.operations.optimizations.utils import load_model
 from torchvision import models
 
-from speedster import optimize_model
+from speedster import optimize_model, load_model
 from speedster.api.tests.utils import torch_to_onnx
 
 
@@ -57,16 +56,20 @@ def test_onnx_ort():
             assert isinstance(loaded_model.get_size(), int)
 
             # Try the optimized model
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            device = torch.device(
+                "cuda" if torch.cuda.is_available() else "cpu"
+            )
             x = torch.randn(1, 3, 256, 256, requires_grad=False)
             model.eval()
             res_original = model(x.to(device))
             res_optimized = optimized_model(x.numpy())[0]
 
             assert (
-                abs((res_original.detach().cpu().numpy() - res_optimized)).max()
+                abs(
+                    (res_original.detach().cpu().numpy() - res_optimized)
+                ).max()
                 < 1e-2
-        )
+            )
 
 
 def test_onnx_ort_quant():

apps/accelerate/speedster/speedster/api/tests/test_pytorch.py

@@ -28,9 +28,8 @@
     tvm_is_available,
     bladedisc_is_available,
 )
-from nebullvm.operations.optimizations.utils import load_model
 
-from speedster import optimize_model
+from speedster import optimize_model, load_model
 
 
 def test_torch_ort():

apps/accelerate/speedster/speedster/api/tests/test_tensorflow.py

@@ -1,6 +1,6 @@
-import cpuinfo
 from tempfile import TemporaryDirectory
 
+import cpuinfo
 import pytest
 import tensorflow as tf
 from keras.applications import ResNet50
@@ -22,10 +22,9 @@
     TensorflowApacheTVMInferenceLearner,
 )
 from nebullvm.operations.optimizations.compilers.utils import tvm_is_available
-from nebullvm.operations.optimizations.utils import load_model
 from nebullvm.tools.utils import gpu_is_available
 
-from speedster import optimize_model
+from speedster import optimize_model, load_model
 
 # Limit tensorflow gpu memory usage
 gpus = tf.config.list_physical_devices("GPU")

apps/accelerate/speedster/speedster/root_op.py

@@ -71,7 +71,7 @@
 SPEEDSTER_FEEDBACK_COLLECTOR = FeedbackCollector(
     url="https://nebuly.cloud/v1/store_speedster_results",
     disable_telemetry_environ_var="SPEEDSTER_DISABLE_TELEMETRY",
-    app_version="0.1.3",
+    app_version="0.2.0",
 )
 
 

nebullvm/config.py

@@ -1,7 +1,7 @@
 from nebullvm.optional_modules.torch import torch
 
 
-VERSION = "0.7.3"
+VERSION = "0.8.0"
 LEARNER_METADATA_FILENAME = "metadata.json"
 ONNX_OPSET_VERSION = 13
 NEBULLVM_DEBUG_FILE = "nebullvm_debug.json"

nebullvm/operations/inference_learners/utils.py

@@ -0,0 +1,12 @@
+from pathlib import Path
+from typing import Union
+
+from nebullvm.operations.inference_learners.base import LearnerMetadata
+
+
+def load_model(path: Union[Path, str]):
+    return LearnerMetadata.read(path).load_model(path)
+
+
+def save_model(model, path: Union[Path, str]):
+    model.save(path)

nebullvm/operations/optimizations/tests/test_deepsparse.py

@@ -19,7 +19,7 @@
     deepsparse_is_available,
 )
 from nebullvm.operations.optimizations.tests.utils import initialize_model
-from nebullvm.operations.optimizations.utils import load_model
+from nebullvm.operations.inference_learners.utils import load_model
 from nebullvm.tools.base import DeepLearningFramework, Device, ModelCompiler
 
 device = Device.CPU

nebullvm/operations/optimizations/tests/test_intel_neural_compressor.py

@@ -19,7 +19,7 @@
     initialize_model,
     check_model_validity,
 )
-from nebullvm.operations.optimizations.utils import load_model
+from nebullvm.operations.inference_learners.utils import load_model
 from nebullvm.tools.base import (
     DeepLearningFramework,
     QuantizationType,

nebullvm/operations/optimizations/tests/test_onnxruntime.py

@@ -18,7 +18,7 @@
     initialize_model,
     check_model_validity,
 )
-from nebullvm.operations.optimizations.utils import load_model
+from nebullvm.operations.inference_learners.utils import load_model
 from nebullvm.tools.base import (
     DeepLearningFramework,
     QuantizationType,

nebullvm/operations/optimizations/tests/test_openvino.py

@@ -19,7 +19,7 @@
     initialize_model,
     check_model_validity,
 )
-from nebullvm.operations.optimizations.utils import load_model
+from nebullvm.operations.inference_learners.utils import load_model
 from nebullvm.tools.base import (
     Device,
     DeepLearningFramework,

nebullvm/operations/optimizations/tests/test_tensor_rt.py

@@ -20,7 +20,7 @@
     initialize_model,
     check_model_validity,
 )
-from nebullvm.operations.optimizations.utils import load_model
+from nebullvm.operations.inference_learners.utils import load_model
 from nebullvm.tools.base import (
     DeepLearningFramework,
     QuantizationType,

nebullvm/operations/optimizations/tests/test_tensorflow.py

@@ -17,7 +17,7 @@
     initialize_model,
     check_model_validity,
 )
-from nebullvm.operations.optimizations.utils import load_model
+from nebullvm.operations.inference_learners.utils import load_model
 from nebullvm.tools.base import (
     DeepLearningFramework,
     QuantizationType,

nebullvm/operations/optimizations/tests/test_torchscript.py

@@ -16,7 +16,7 @@
     initialize_model,
     check_model_validity,
 )
-from nebullvm.operations.optimizations.utils import load_model
+from nebullvm.operations.inference_learners.utils import load_model
 from nebullvm.tools.base import (
     DeepLearningFramework,
     QuantizationType,

nebullvm/operations/optimizations/tests/test_tvm.py

@@ -20,7 +20,7 @@
     initialize_model,
     check_model_validity,
 )
-from nebullvm.operations.optimizations.utils import load_model
+from nebullvm.operations.inference_learners.utils import load_model
 from nebullvm.tools.base import (
     DeepLearningFramework,
     QuantizationType,

nebullvm/operations/optimizations/utils.py

@@ -1,14 +1,9 @@
-from pathlib import Path
-from typing import Callable, List, Union
+from typing import Callable, List
 
-from nebullvm.operations.inference_learners.base import LearnerMetadata
 
 def map_compilers_and_compressors(ignore_list: List, enum_class: Callable):
     if ignore_list is None:
         ignore_list = []
     else:
         ignore_list = [enum_class(element) for element in ignore_list]
     return ignore_list
-
-def load_model(path: Union[Path, str]):
-    return LearnerMetadata.read(path).load_model(path)

notebooks/speedster/huggingface/Accelerate_Hugging_Face_PyTorch_BERT_with_Speedster.ipynb

@@ -222,7 +222,7 @@
       },
       "outputs": [],
       "source": [
-        "from speedster import optimize_model"
+        "from speedster import optimize_model, save_model, load_model"
       ]
     },
     {
@@ -572,7 +572,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "optimized_model.save(\"model_save_path\")"
+        "save_model(optimized_model, \"model_save_path\")"
       ]
     },
     {
@@ -591,9 +591,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "from nebullvm.operations.inference_learners.base import LearnerMetadata\n",
-        "\n",
-        "optimized_model = LearnerMetadata.read(\"model_save_path\").load_model(\"model_save_path\")"
+        "optimized_model = load_model(\"model_save_path\")"
       ]
     },
     {

notebooks/speedster/huggingface/Accelerate_Hugging_Face_PyTorch_DistilBERT_with_Speedster.ipynb

@@ -222,7 +222,7 @@
       },
       "outputs": [],
       "source": [
-        "from speedster import optimize_model"
+        "from speedster import optimize_model, save_model, load_model"
       ]
     },
     {
@@ -572,7 +572,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "optimized_model.save(\"model_save_path\")"
+        "save_model(optimized_model, \"model_save_path\")"
       ]
     },
     {
@@ -591,9 +591,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "from nebullvm.operations.inference_learners.base import LearnerMetadata\n",
-        "\n",
-        "optimized_model = LearnerMetadata.read(\"model_save_path\").load_model(\"model_save_path\")"
+        "optimized_model = load_model(\"model_save_path\")"
       ]
     },
     {

notebooks/speedster/huggingface/Accelerate_Hugging_Face_PyTorch_GPT2_with_Speedster.ipynb

@@ -231,7 +231,7 @@
       },
       "outputs": [],
       "source": [
-        "from speedster import optimize_model"
+        "from speedster import optimize_model, save_model, load_model"
       ]
     },
     {
@@ -545,7 +545,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "optimized_model.save(\"model_save_path\")"
+        "save_model(optimized_model, \"model_save_path\")"
       ]
     },
     {
@@ -564,9 +564,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "from nebullvm.operations.inference_learners.base import LearnerMetadata\n",
-        "\n",
-        "optimized_model = LearnerMetadata.read(\"model_save_path\").load_model(\"model_save_path\")"
+        "optimized_model = load_model(\"model_save_path\")"
       ]
     },
     {

notebooks/speedster/huggingface/Accelerate_Hugging_Face_TensorFlow_BERT_with_Speedster.ipynb

@@ -217,7 +217,7 @@
       },
       "outputs": [],
       "source": [
-        "from speedster import optimize_model"
+        "from speedster import optimize_model, save_model, load_model"
       ]
     },
     {
@@ -560,7 +560,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "optimized_model.save(\"model_save_path\")"
+        "save_model(optimized_model, \"model_save_path\")"
       ]
     },
     {
@@ -579,9 +579,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "from nebullvm.operations.inference_learners.base import LearnerMetadata\n",
-        "\n",
-        "optimized_model = LearnerMetadata.read(\"model_save_path\").load_model(\"model_save_path\")"
+        "optimized_model = load_model(\"model_save_path\")"
       ]
     },
     {

notebooks/speedster/onnx/Accelerate_ONNX_ResNet50_with_Speedster.ipynb

@@ -171,7 +171,7 @@
       "outputs": [],
       "source": [
         "import numpy as np\n",
-        "from speedster import optimize_model\n",
+        "from speedster import optimize_model, save_model, load_model\n",
         "\n",
         "# Load a resnet as example\n",
         "model = \"resnet50-v1-12.onnx\"\n",
@@ -540,7 +540,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "optimized_model.save(\"model_save_path\")"
+        "save_model(optimized_model, \"model_save_path\")"
       ]
     },
     {
@@ -559,9 +559,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "from nebullvm.operations.inference_learners.base import LearnerMetadata\n",
-        "\n",
-        "optimized_model = LearnerMetadata.read(\"model_save_path\").load_model(\"model_save_path\")"
+        "optimized_model = load_model(\"model_save_path\")"
       ]
     },
     {