[SD-126] Save model summary during benchmarking

jetson/power_logging/model/save_model_summary.py

@@ -0,0 +1,102 @@
+import torch
+import json
+from typing import Any 
+# from torchsummary import summary
+from torchinfo import summary
+# 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, pretrained=True)
+    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 get_layer_info(model, input_shape):
+    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__,
+                # "kernal_size": layer.kernel_size if hasattr(layer, "kernel_size") else None,
+                # "stride": layer.stride if hasattr(layer, "stride") else None,
+                # "padding": layer.padding if hasattr(layer, "padding") else None,
+                # "type": type(layer)
+            }
+        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
+
+
+model = load_model("resnet18", "pytorch/vision:v0.10.0")
+layer_info = get_layer_info(model, (1, 3, 224, 224))
+# print(len(layer_info))
+
+print(json.dumps(layer_info, indent=4, separators=(",", ": "), ensure_ascii=False, default=str))
+
+
+# save to json later

jetson/power_logging/pyproject.toml

@@ -9,6 +9,8 @@ dependencies = [
     "pandas>=2.2.3",
     "pillow>=11.1.0",
     "torch==2.2.2",
+    "torchinfo>=1.8.0",
+    "torchsummary>=1.5.1",
 ]
 
 [dependency-groups]