diff --git a/ecoml/src/ecoml/ecoml.py b/ecoml/src/ecoml/ecoml.py index cc3bc8d..e9b65e8 100644 --- a/ecoml/src/ecoml/ecoml.py +++ b/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 diff --git a/ecoml/src/ecoml/infer.py b/ecoml/src/ecoml/infer.py index ec29514..0c4a004 100644 --- a/ecoml/src/ecoml/infer.py +++ b/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 diff --git a/ecoml/src/ecoml/model_summary/model_summary.py b/ecoml/src/ecoml/model_summary/model_summary.py index 6de6592..70fc682 100644 --- a/ecoml/src/ecoml/model_summary/model_summary.py +++ b/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)