[SD-161] CLI PyTorch Integration

ecoml/src/ecoml/ecoml.py

@@ -3,10 +3,12 @@
 import json
 from pathlib import Path
 from typing import Annotated
-
+import torch
 import typer
 from rich.console import Console
 from rich.table import Table
+from ecoml.model_summary.model_summary import get_summary
+from ecoml.infer import run_inference, display_latency_table, display_metrics_table, display_runtime_table
 
 CONFIG = {"jetson_orin": {"pytorch": {"low": 5, "average": 7, "high": 10}}}
 
@@ -34,6 +36,14 @@ def validate_model(model_path: str):
             return False, None
     return False, None
 
+def load_model(model_path: str):
+    suffix = Path(model_path).suffix
+    if suffix in [".pth", ".pt"]:
+        model = torch.load(model_path, map_location=torch.device('cpu'))
+        if isinstance(model, dict) and "model" in model:
+            model = model["model"]
+        return model
+    return None
 
 def display_config_table(cfg: dict[str, int]) -> None:
     """Display a table of power profiles.
@@ -70,26 +80,40 @@ def predict(
     If --verbose is used, a detailed summary of predictions is provided.
     """
     cfg = CONFIG["jetson_orin"]["pytorch"]
-
-    success, _ = validate_model(model)
-    if not success:
-        error_console.print("Expected a valid PyTorch model summary JSON File.")
+    model_path = Path(model)
+
+    if model_path.suffix in [".pth", ".pt"]:
+        pytorch_model = load_model(model)
+        if pytorch_model is None:
+            error_console.print("Failed to load model...")
+            raise typer.Exit(code=1)
+
+        summary = get_summary(pytorch_model)
+    elif model_path.suffix == ".json":
+        summary = model_path
+    else:
+        error_console.print("Invalid file type...")
         raise typer.Exit(code=1)
+
+    # success, _ = validate_model(model)
+    # if not success:
+    #     error_console.print("Expected a valid PyTorch model summary JSON File.")
+    #     raise typer.Exit(code=1)
 
-    # Import relevant functions
-    from ecoml.infer import(
-        run_inference,
-        display_latency_table,
-        display_metrics_table,
-        display_runtime_table
-    )
+    # # Import relevant functions
+    # from ecoml.infer import(
+    #     run_inference,
+    #     display_latency_table,
+    #     display_metrics_table,
+    #     display_runtime_table
+    # )
 
     # Run the inference function that returns the dictionary
-    runtime_predictions = run_inference(Path(model), power_profiles=cfg, verbose=verbose)
+    runtime_predictions = run_inference(summary, power_profiles=cfg, verbose=verbose)
 
     # If it is an empty dict then throw an error
     if not runtime_predictions:
-        error_console.print("Inference failed. No results were returned")
+        error_console.print("Inference failed. No results were returned...")
         raise typer.Exit(code=1)
 
     # Take out data from the dict

ecoml/src/ecoml/infer.py

@@ -9,9 +9,11 @@
 from rich.console import Console
 from rich.table import Table
 from dataclasses import dataclass
+from typing import Union
 
-from ecoml.data_preparation.pytorch_utils import read_layers_info
+from ecoml.data_preparation.pytorch_utils import PytorchLayer, read_layers_info
 from ecoml.model_builder.model_inference import InferenceModel
+from ecoml.model_summary.model_summary import get_summary
 
 console = Console()
 error_console = Console(stderr=True, style="bold red")
@@ -22,19 +24,25 @@ class InferenceResult:
     ltype: str
     runtime: float
 
-def run_inference(model_sumary_path: Path, power_profiles: dict[str, int], verbose: bool = False) -> list[InferenceResult]:
+def run_inference(model_input: Union[Path, dict], power_profiles: dict[str, int], verbose: bool = False) -> list[InferenceResult]:
     convolution = InferenceModel(model_version=1, layer_type="convolutional", verbose=verbose)
     pooling = InferenceModel(model_version=1, layer_type="pooling", verbose=verbose)
     dense = InferenceModel(model_version=1, layer_type="dense", verbose=verbose)
 
-    layer_info_read = read_layers_info(model_sumary_path)
+    if isinstance(model_input, Path):
+        layer_info_read = read_layers_info(model_input)
+    elif isinstance(model_input, dict):
+        layer_info_read = {k: PytorchLayer(**v) for k, v in model_input.items()}
+    else:
+        error_console.print("Invalid model input...")
+        return []
 
     if verbose:
-        print(f"Found {len(layer_info_read)} layers in {model_sumary_path}.")
+        print(f"Found {len(layer_info_read)} layers in model...")
 
     inference_results = []
     for layer_name, layer_info in layer_info_read.items():
-        layer_type = layer_info.get_layer_type()
+        layer_type = layer_info.get_layer_type() if hasattr(layer_info, 'get_layer_type') else layer_info.get('type', 'unknown')
 
         if layer_type == "convolutional":
             model = convolution

ecoml/src/ecoml/model_summary/model_summary.py

@@ -4,6 +4,9 @@
 from pathlib import Path
 
 import torch
+import torch.nn.quantized as quantized_nn
+from torch.nn.intrinsic.quantized import ConvReLU2d
+
 
 
 def get_layers(
@@ -19,18 +22,32 @@ def get_layers(
     Returns:
         a list of tuple containing the layer name and the layer.
     """
-    children = list(model.named_children())
+    if not hasattr(model, "_modules") or isinstance(model, (quantized_nn.Conv2d, ConvReLU2d, quantized_nn.Linear, quantized_nn.BatchNorm2d)):
+        return []
+
+    try:
+        children = list(model.named_children())
+    except AttributeError:
+        return [(name_prefix, model)]
 
     if len(children) == 0:
-        result = [(name_prefix, model)]
-    else:
-        result = []
-        for child_name, child in children:
-            layers = get_layers(child, name_prefix + "_" + child_name)
-            result.extend(layers)
+        return [(name_prefix, model)]
+
+    result = []
+    for child_name, child in children:
+        layers = get_layers(child, name_prefix + "_" + child_name)
+        result.extend(layers)
 
     return result
 
+def is_quantized_model(model: torch.nn.Module) -> bool:
+    try:
+        return any(
+            isinstance(layer, (quantized_nn.Conv2d, ConvReLU2d, quantized_nn.Linear))
+            for _, layer in model.named_modules()
+        )
+    except AttributeError:
+        return False
 
 def get_summary(
     model: torch.nn.Module,
@@ -52,6 +69,23 @@ def get_summary(
     test = torch.randn(*input_shape)
     hooks = []
 
+    if is_quantized_model(model):
+        print("Detected quantized model...")
+        for layer in model.children():
+            try:
+                model_info[layer.__class__.__name__] = {
+                    "type": layer.__class__.__name__,
+                    "kernel_size": getattr(layer, "kernel_size", None),
+                    "stride": getattr(layer, "stride", None),
+                    "padding": getattr(layer, "padding", None),
+                    "input_shape": [1, 3, 224, 224],
+                    "output_shape": [1, 3, 224, 224],
+                }
+            except AttributeError:
+                print(f"Skipping layer {layer.__class__.__name__} as it misses attribute...")
+                continue
+        return model_info
+
     def register_hook(layer_name):
         def hook(module, input, output):
             model_info[layer_name] = {
@@ -65,10 +99,22 @@ def hook(module, input, output):
 
         return hook
 
+    valid_layers = []
     for layer_name, layer in get_layers(model):
+        if not hasattr(layer, "register_forward_hook"):
+            continue
+        valid_layers.append((layer_name, layer))
+
+    if not valid_layers:
+        raise ValueError("No valid layers found...")
+
+    for layer_name, layer in valid_layers:
         hooks.append(layer.register_forward_hook(register_hook(layer_name)))
 
-    model.eval()
+    try:
+        model.eval()
+    except AttributeError:
+        print("model.eval() could not  be applied")
     with torch.no_grad():
         _ = model(test)