[SD-126] Save model summary during benchmarking

jetson/power_logging/model/benchmark.py

@@ -18,6 +18,7 @@
 
 from model.lenet import LeNet
 from functools import partial
+from model_utils import load_model, get_layers
 
 """
 Wrapper class for Torch.cuda.event for non-CUDA supported devices
@@ -68,60 +69,6 @@ class BenchmarkMetrics(BaseModel):
     avg_latency: float  # in seconds
     avg_throughput: float
 
-
-def load_model(model_name: str, model_repo: str) -> Any:
-    """Load model from Pytorch Hub.
-
-    Args:
-        model_name: Name of model.
-            It should be same as that in Pytorch Hub.
-
-    Raises:
-        ValueError: If loading model fails from PyTorch Hub
-
-    Returns:
-        PyTorch model
-    """
-    if model_name == "lenet":
-        return LeNet()
-    if model_name == "fcn_resnet50":
-        return torch.hub.load(model_repo, model_name, pretrained=True)
-    try:
-        return torch.hub.load(model_repo, model_name)
-    except:
-        raise ValueError(
-            f"Model name: {model_name} is most likely incorrect. "
-            "Please refer https://pytorch.org/hub/ to get model name."
-        )
-
-
-def get_layers(model: torch.nn.Module, name_prefix: str="") -> list[tuple[str, torch.nn.Module]]:
-    """
-    Recursively get all layers in a pytorch model.
-
-    Args:
-        model: the pytorch model to look for layers.
-        name_prefix: Use to identify the parents layer. Defaults to "".
-
-    Returns:
-        a list of tuple containing the layer name and the layer.
-    """
-    children = list(model.named_children())
-
-    if len(children) == 0: # No child
-        result = [(name_prefix, model)]
-    else:
-        # If have children, iterate over each child.
-        result = []
-        for child_name, child in children:
-            # Recursively call get_layers on the child, appending the current
-            # child's name to the name_prefix.
-            layers = get_layers(child, name_prefix + "_" + child_name)
-            result.extend(layers)
-
-    return result
-
-
 def define_and_register_hooks(model, device) -> dict:
     """
         Define and register hooks with CUDA or CPU timing.

jetson/power_logging/model/generate_model_summaries.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+
+# Directory to store results
+RESULT_DIR = "raw_data/prebuilt_models"
+
+# List of models to process
+MODELS = ("alexnet" "vgg11" "vgg13" "vgg16" "vgg19" "mobilenet_v2" "mobilenet_v3_small" "mobilenet_v3_large" "resnet18" "resnet34" "resnet50" "resnet101" "resnet152" "lenet" "resnext50_32x4d" "resnext101_32x8d" "resnext101_64x4d" "convnext_tiny" "convnext_small" "convnext_base")
+
+PYTHON_SCRIPT = "save_model_summary.py"
+
+for MODEL in "${MODELS[@]}"
+do
+    MODEL_OUTPUT_DIR="${RESULT_DIR}/${MODEL}"
+    OUTPUT_FILE = "${MODEL_OUTPUT_DIR}/model_summary.json"
+
+    mkdir -p "$MODEL_OUTPUT_DIR"
+
+    echo "Generating summary for model: $MODEL"
+    python "$PYTHON_SCRIPT" --model "$MODEL" --output-file "$OUTPUT_FILE"
+
+    if [ $? -eq 0]; then
+        echo "Summary saved to: $OUTPUT_FILE"
+    else
+        echo "Failed to generate summary for model: $MODEL"
+    fi
+done

jetson/power_logging/model/model_utils.py

@@ -0,0 +1,51 @@
+import torch
+from typing import Any
+from model.lenet import LeNet
+
+def load_model(model_name: str, model_repo: str) -> Any:
+    """Load model from Pytorch Hub.
+
+    Args:
+        model_name: Name of model.
+            It should be same as that in Pytorch Hub.
+
+    Raises:
+        ValueError: If loading model fails from PyTorch Hub
+
+    Returns:
+        PyTorch model
+    """
+    if model_name == "lenet":
+        return LeNet()
+    if model_name == "fcn_resnet50":
+        return torch.hub.load(model_repo, model_name, pretrained=True)
+    try:
+        return torch.hub.load(model_repo, model_name)
+    except:
+        raise ValueError(
+            f"Model name: {model_name} is most likely incorrect. "
+            "Please refer https://pytorch.org/hub/ to get model name."
+        )
+
+def get_layers(model: torch.nn.Module, name_prefix: str="") -> list[tuple[str, torch.nn.Module]]:
+    """
+    Recursively get all layers in a pytorch model.
+
+    Args:
+        model: the pytorch model to look for layers.
+        name_prefix: Use to identify the parents layer. Defaults to "".
+
+    Returns:
+        a list of tuple containing the layer name and the layer.
+    """
+    children = list(model.named_children())
+
+    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 result

jetson/power_logging/model/save_model_summary.py

@@ -0,0 +1,85 @@
+import torch
+import json
+import argparse
+from typing import Any 
+from model_utils import load_model, get_layers
+
+def get_layer_info(model, input_shape):
+    """
+    Get key information of all layers within a model.
+
+    Args:
+        model: The pytorch model
+        input_shape: input size of the model
+
+    Returns:
+        information about the model
+    """
+    model_info = {}
+    test = torch.randn(*input_shape)
+    hooks = []
+
+    def register_hook(layer_name):
+        def hook(module, input, output):
+            model_info[layer_name] = {
+                "input_shape": tuple(input[0].size()) if input else None,
+                "output_shape": tuple(output.size()) if output is not None else None,
+                "kernel_size": getattr(module, "kernel_size", None),
+                "stride": getattr(module, "stride", None),
+                "padding": getattr(module, "padding", None),
+                "type": module.__class__.__name__,
+            }
+        return hook
+
+    for layer_name, layer in get_layers(model):
+        hooks.append(layer.register_forward_hook(register_hook(layer_name)))
+
+    model.eval()
+    with torch.no_grad():
+        _ = model(test)
+
+    for hook in hooks:
+        hook.remove()
+
+    return model_info
+
+
+def run(args):
+    model = load_model(args.model, args.model_repo)
+    layer_info = get_layer_info(model, args.input_shape)
+
+    print(json.dumps(layer_info, indent=4, separators=(",", ": "), ensure_ascii=False))
+
+    output = "model_summary.json"
+    with open(output, "w") as file:
+        json.dump(layer_info, file, indent=4, separators=(",", ": "), ensure_ascii=False)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        prog="Save Model Summary",
+        description="Save a summary of the model after benchmarking.",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="resnet18",
+        help="Specify name of pretrained CNN mode from PyTorch Hub."
+        "For more information on PyTorch Hub visit: "
+        "https://pytorch.org/hub/research-models",
+    )
+    parser.add_argument(
+        "--model-repo",
+        type=str,
+        default="pytorch/vision:v0.10.0",
+        help="Specify path and version to model repository from PyTorch Hub."
+    )
+    parser.add_argument(
+        "--input-shape",
+        type=int,
+        nargs="+",
+        default=[1, 3, 224, 224],
+        help="Input shape BCHW",
+    )
+
+    args = parser.parse_args()
+    run(args)