""" PyTorch Model Analyzer Extracts architecture information from PyTorch models for 3D visualization. """ import torch import torch.nn as nn from typing import Dict, List, Any, Optional, Tuple from dataclasses import dataclass, asdict from collections import OrderedDict import json @dataclass class LayerInfo: """Information about a single layer in the model.""" id: str name: str type: str category: str input_shape: Optional[List[int]] output_shape: Optional[List[int]] params: Dict[str, Any] num_parameters: int trainable: bool @dataclass class ConnectionInfo: """Information about connections between layers.""" source: str target: str tensor_shape: Optional[List[int]] @dataclass class ModelArchitecture: """Complete model architecture information.""" name: str framework: str total_parameters: int trainable_parameters: int layers: List[LayerInfo] connections: List[ConnectionInfo] input_shape: Optional[List[int]] output_shape: Optional[List[int]] # Layer category mapping LAYER_CATEGORIES = { # Convolution layers 'Conv1d': 'convolution', 'Conv2d': 'convolution', 'Conv3d': 'convolution', 'ConvTranspose1d': 'convolution', 'ConvTranspose2d': 'convolution', 'ConvTranspose3d': 'convolution', # Pooling layers 'MaxPool1d': 'pooling', 'MaxPool2d': 'pooling', 'MaxPool3d': 'pooling', 'AvgPool1d': 'pooling', 'AvgPool2d': 'pooling', 'AvgPool3d': 'pooling', 'AdaptiveAvgPool1d': 'pooling', 'AdaptiveAvgPool2d': 'pooling', 'AdaptiveAvgPool3d': 'pooling', 'AdaptiveMaxPool1d': 'pooling', 'AdaptiveMaxPool2d': 'pooling', 'AdaptiveMaxPool3d': 'pooling', 'GlobalAveragePooling2D': 'pooling', # Linear/Dense layers 'Linear': 'linear', 'LazyLinear': 'linear', 'Bilinear': 'linear', # Normalization layers 'BatchNorm1d': 'normalization', 'BatchNorm2d': 'normalization', 'BatchNorm3d': 'normalization', 'LayerNorm': 'normalization', 'GroupNorm': 'normalization', 'InstanceNorm1d': 'normalization', 'InstanceNorm2d': 'normalization', 'InstanceNorm3d': 'normalization', # Activation layers 'ReLU': 'activation', 'ReLU6': 'activation', 'LeakyReLU': 'activation', 'PReLU': 'activation', 'ELU': 'activation', 'SELU': 'activation', 'GELU': 'activation', 'Sigmoid': 'activation', 'Tanh': 'activation', 'Softmax': 'activation', 'LogSoftmax': 'activation', 'Softplus': 'activation', 'Softsign': 'activation', 'Hardswish': 'activation', 'Hardsigmoid': 'activation', 'SiLU': 'activation', 'Mish': 'activation', # Dropout layers 'Dropout': 'regularization', 'Dropout2d': 'regularization', 'Dropout3d': 'regularization', 'AlphaDropout': 'regularization', # Recurrent layers 'RNN': 'recurrent', 'LSTM': 'recurrent', 'GRU': 'recurrent', 'RNNCell': 'recurrent', 'LSTMCell': 'recurrent', 'GRUCell': 'recurrent', # Transformer layers 'Transformer': 'attention', 'TransformerEncoder': 'attention', 'TransformerDecoder': 'attention', 'TransformerEncoderLayer': 'attention', 'TransformerDecoderLayer': 'attention', 'MultiheadAttention': 'attention', # Embedding layers 'Embedding': 'embedding', 'EmbeddingBag': 'embedding', # Reshape/View layers 'Flatten': 'reshape', 'Unflatten': 'reshape', # Container layers 'Sequential': 'container', 'ModuleList': 'container', 'ModuleDict': 'container', } def get_layer_category(layer_type: str) -> str: """Get the category for a layer type.""" return LAYER_CATEGORIES.get(layer_type, 'other') def count_parameters(module: nn.Module) -> Tuple[int, int]: """Count total and trainable parameters in a module.""" total = sum(p.numel() for p in module.parameters()) trainable = sum(p.numel() for p in module.parameters() if p.requires_grad) return total, trainable def extract_layer_params(module: nn.Module, layer_type: str) -> Dict[str, Any]: """Extract relevant parameters from a layer.""" params = {} try: if hasattr(module, 'in_features'): params['in_features'] = module.in_features if hasattr(module, 'out_features'): params['out_features'] = module.out_features if hasattr(module, 'in_channels'): params['in_channels'] = module.in_channels if hasattr(module, 'out_channels'): params['out_channels'] = module.out_channels if hasattr(module, 'kernel_size'): ks = module.kernel_size params['kernel_size'] = list(ks) if isinstance(ks, tuple) else ks if hasattr(module, 'stride'): s = module.stride params['stride'] = list(s) if isinstance(s, tuple) else s if hasattr(module, 'padding'): p = module.padding params['padding'] = list(p) if isinstance(p, tuple) else p if hasattr(module, 'dilation'): d = module.dilation params['dilation'] = list(d) if isinstance(d, tuple) else d if hasattr(module, 'groups'): params['groups'] = module.groups if hasattr(module, 'bias') and module.bias is not None: params['bias'] = True if hasattr(module, 'num_features'): params['num_features'] = module.num_features if hasattr(module, 'eps'): params['eps'] = module.eps if hasattr(module, 'momentum') and module.momentum is not None: params['momentum'] = module.momentum if hasattr(module, 'normalized_shape'): params['normalized_shape'] = list(module.normalized_shape) if hasattr(module, 'hidden_size'): params['hidden_size'] = module.hidden_size if hasattr(module, 'num_layers'): params['num_layers'] = module.num_layers if hasattr(module, 'bidirectional'): params['bidirectional'] = module.bidirectional if hasattr(module, 'num_heads'): params['num_heads'] = module.num_heads if hasattr(module, 'embed_dim'): params['embed_dim'] = module.embed_dim if hasattr(module, 'num_embeddings'): params['num_embeddings'] = module.num_embeddings if hasattr(module, 'embedding_dim'): params['embedding_dim'] = module.embedding_dim if hasattr(module, 'p') and layer_type.startswith('Dropout'): params['p'] = module.p if hasattr(module, 'negative_slope'): params['negative_slope'] = module.negative_slope if hasattr(module, 'inplace'): params['inplace'] = module.inplace if hasattr(module, 'dim'): params['dim'] = module.dim except Exception: pass return params def analyze_model_structure(model: nn.Module, model_name: str = "model") -> ModelArchitecture: """ Analyze a PyTorch model and extract its architecture. Args: model: The PyTorch model to analyze model_name: Name identifier for the model Returns: ModelArchitecture with complete layer and connection information """ layers = [] connections = [] layer_index = 0 parent_stack = [] def process_module(name: str, module: nn.Module, parent_id: Optional[str] = None): nonlocal layer_index layer_type = module.__class__.__name__ # Skip container modules but process their children if layer_type in ('Sequential', 'ModuleList', 'ModuleDict'): for child_name, child in module.named_children(): full_name = f"{name}.{child_name}" if name else child_name process_module(full_name, child, parent_id) return # Skip modules with no parameters and no meaningful operation # But include activation, pooling, dropout, etc. has_params = sum(1 for _ in module.parameters(recurse=False)) > 0 is_meaningful = layer_type in LAYER_CATEGORIES or has_params if not is_meaningful and len(list(module.children())) > 0: # Process children of non-meaningful containers for child_name, child in module.named_children(): full_name = f"{name}.{child_name}" if name else child_name process_module(full_name, child, parent_id) return layer_id = f"layer_{layer_index}" layer_index += 1 total_params, trainable_params = count_parameters(module) params = extract_layer_params(module, layer_type) layer_info = LayerInfo( id=layer_id, name=name or layer_type, type=layer_type, category=get_layer_category(layer_type), input_shape=None, # Will be populated during forward pass output_shape=None, params=params, num_parameters=total_params, trainable=trainable_params > 0 ) layers.append(layer_info) # Create connection from parent if parent_id is not None: connections.append(ConnectionInfo( source=parent_id, target=layer_id, tensor_shape=None )) # Process children children = list(module.named_children()) if children: for child_name, child in children: full_name = f"{name}.{child_name}" if name else child_name process_module(full_name, child, layer_id) return layer_id # Process the model children = list(model.named_children()) if children: prev_id = None for name, child in children: layer_id = process_module(name, child, prev_id) if layer_id: prev_id = layer_id else: # Single layer model process_module("", model, None) # If layers are sequential and no connections exist, create linear connections if len(layers) > 1 and len(connections) == 0: for i in range(len(layers) - 1): connections.append(ConnectionInfo( source=layers[i].id, target=layers[i + 1].id, tensor_shape=None )) total_params, trainable_params = count_parameters(model) return ModelArchitecture( name=model_name, framework="pytorch", total_parameters=total_params, trainable_parameters=trainable_params, layers=layers, connections=connections, input_shape=None, output_shape=None ) def trace_model_shapes(model: nn.Module, input_tensor: torch.Tensor, arch: ModelArchitecture) -> ModelArchitecture: """ Trace model execution to capture input/output shapes for each layer. Args: model: The PyTorch model input_tensor: Sample input tensor arch: Existing architecture info to update Returns: Updated ModelArchitecture with shape information """ shapes = {} hooks = [] def make_hook(name): def hook(module, input, output): input_shape = None output_shape = None if isinstance(input, tuple) and len(input) > 0: if isinstance(input[0], torch.Tensor): input_shape = list(input[0].shape) elif isinstance(input, torch.Tensor): input_shape = list(input.shape) if isinstance(output, torch.Tensor): output_shape = list(output.shape) elif isinstance(output, tuple) and len(output) > 0: if isinstance(output[0], torch.Tensor): output_shape = list(output[0].shape) shapes[name] = { 'input': input_shape, 'output': output_shape } return hook # Register hooks for name, module in model.named_modules(): if name: # Skip root module hooks.append(module.register_forward_hook(make_hook(name))) # Run forward pass try: model.eval() with torch.no_grad(): output = model(input_tensor) # Update architecture with shapes for layer in arch.layers: if layer.name in shapes: layer.input_shape = shapes[layer.name]['input'] layer.output_shape = shapes[layer.name]['output'] # Set model input/output shapes arch.input_shape = list(input_tensor.shape) if isinstance(output, torch.Tensor): arch.output_shape = list(output.shape) except Exception as e: print(f"Warning: Could not trace shapes: {e}") finally: # Remove hooks for hook in hooks: hook.remove() return arch def load_pytorch_model(file_path: str) -> Tuple[Optional[nn.Module], Optional[Dict], str]: """ Load a PyTorch model from file. Returns: Tuple of (model, state_dict, model_type) model_type can be: 'full_model', 'state_dict', 'torchscript', 'checkpoint' """ try: # Try loading as TorchScript first try: model = torch.jit.load(file_path, map_location='cpu') return model, None, 'torchscript' except Exception: pass # Try loading as regular checkpoint checkpoint = torch.load(file_path, map_location='cpu', weights_only=False) if isinstance(checkpoint, nn.Module): return checkpoint, None, 'full_model' if isinstance(checkpoint, dict): # Check for common checkpoint formats if 'model' in checkpoint: if isinstance(checkpoint['model'], nn.Module): return checkpoint['model'], None, 'checkpoint' elif isinstance(checkpoint['model'], dict): return None, checkpoint['model'], 'state_dict' if 'state_dict' in checkpoint: return None, checkpoint['state_dict'], 'state_dict' if 'model_state_dict' in checkpoint: return None, checkpoint['model_state_dict'], 'state_dict' # Check if it's directly a state dict (contains tensor values) has_tensors = any(isinstance(v, torch.Tensor) for v in checkpoint.values()) if has_tensors: return None, checkpoint, 'state_dict' return None, None, 'unknown' except Exception as e: raise ValueError(f"Failed to load model: {str(e)}") def analyze_state_dict(state_dict: Dict[str, torch.Tensor], model_name: str = "model") -> ModelArchitecture: """ Analyze a state dict to infer model architecture. This extracts layer information from weight tensor names and shapes. """ layers = [] layer_map = OrderedDict() # Group parameters by layer name for key, tensor in state_dict.items(): if not isinstance(tensor, torch.Tensor): continue # Extract layer name from parameter name parts = key.rsplit('.', 1) if len(parts) == 2: layer_name, param_type = parts else: layer_name = key param_type = 'weight' if layer_name not in layer_map: layer_map[layer_name] = { 'params': {}, 'shapes': {} } layer_map[layer_name]['params'][param_type] = True layer_map[layer_name]['shapes'][param_type] = list(tensor.shape) # Create layer info from grouped parameters layer_index = 0 for layer_name, info in layer_map.items(): layer_type, category = infer_layer_type(layer_name, info['shapes']) layer_id = f"layer_{layer_index}" layer_index += 1 # Compute number of parameters num_params = 0 for param_type, shape in info['shapes'].items(): param_size = 1 for dim in shape: param_size *= dim num_params += param_size # Extract layer parameters from shapes params = extract_params_from_shapes(layer_type, info['shapes']) # Infer input/output shapes input_shape, output_shape = infer_shapes(layer_type, info['shapes'], params) layers.append(LayerInfo( id=layer_id, name=layer_name, type=layer_type, category=category, input_shape=input_shape, output_shape=output_shape, params=params, num_parameters=num_params, trainable=True )) # Create sequential connections connections = [] for i in range(len(layers) - 1): connections.append(ConnectionInfo( source=layers[i].id, target=layers[i + 1].id, tensor_shape=layers[i].output_shape )) total_params = sum(layer.num_parameters for layer in layers) return ModelArchitecture( name=model_name, framework="pytorch", total_parameters=total_params, trainable_parameters=total_params, layers=layers, connections=connections, input_shape=layers[0].input_shape if layers else None, output_shape=layers[-1].output_shape if layers else None ) def infer_layer_type(layer_name: str, shapes: Dict[str, List[int]]) -> Tuple[str, str]: """Infer layer type from name and weight shapes.""" name_lower = layer_name.lower() # Check for common layer type patterns in name if 'conv' in name_lower: weight_shape = shapes.get('weight', []) if len(weight_shape) == 5: return 'Conv3d', 'convolution' elif len(weight_shape) == 4: return 'Conv2d', 'convolution' elif len(weight_shape) == 3: return 'Conv1d', 'convolution' return 'Conv2d', 'convolution' if 'bn' in name_lower or 'batch' in name_lower or 'norm' in name_lower: weight_shape = shapes.get('weight', shapes.get('running_mean', [])) if 'layer' in name_lower: return 'LayerNorm', 'normalization' return 'BatchNorm2d', 'normalization' if 'fc' in name_lower or 'linear' in name_lower or 'dense' in name_lower or 'classifier' in name_lower: return 'Linear', 'linear' if 'lstm' in name_lower: return 'LSTM', 'recurrent' if 'gru' in name_lower: return 'GRU', 'recurrent' if 'rnn' in name_lower: return 'RNN', 'recurrent' if 'attention' in name_lower or 'attn' in name_lower: return 'MultiheadAttention', 'attention' if 'embed' in name_lower: return 'Embedding', 'embedding' if 'pool' in name_lower: return 'AdaptiveAvgPool2d', 'pooling' # Infer from weight shape weight_shape = shapes.get('weight', []) if len(weight_shape) == 2: return 'Linear', 'linear' elif len(weight_shape) == 4: return 'Conv2d', 'convolution' elif len(weight_shape) == 3: return 'Conv1d', 'convolution' elif len(weight_shape) == 1: return 'BatchNorm2d', 'normalization' return 'Unknown', 'other' def extract_params_from_shapes(layer_type: str, shapes: Dict[str, List[int]]) -> Dict[str, Any]: """Extract layer parameters from weight shapes.""" params = {} weight_shape = shapes.get('weight', []) if layer_type in ('Linear',): if len(weight_shape) >= 2: params['out_features'] = weight_shape[0] params['in_features'] = weight_shape[1] params['bias'] = 'bias' in shapes elif layer_type in ('Conv1d', 'Conv2d', 'Conv3d'): if len(weight_shape) >= 2: params['out_channels'] = weight_shape[0] params['in_channels'] = weight_shape[1] if len(weight_shape) > 2: params['kernel_size'] = weight_shape[2:] params['bias'] = 'bias' in shapes elif layer_type in ('BatchNorm1d', 'BatchNorm2d', 'BatchNorm3d'): if len(weight_shape) >= 1: params['num_features'] = weight_shape[0] elif layer_type == 'LayerNorm': if len(weight_shape) >= 1: params['normalized_shape'] = weight_shape elif layer_type == 'Embedding': if len(weight_shape) >= 2: params['num_embeddings'] = weight_shape[0] params['embedding_dim'] = weight_shape[1] elif layer_type in ('LSTM', 'GRU', 'RNN'): # weight_ih_l0 shape gives hidden_size x input_size if 'weight_ih_l0' in shapes: ih_shape = shapes['weight_ih_l0'] if len(ih_shape) >= 2: multiplier = 4 if layer_type == 'LSTM' else (3 if layer_type == 'GRU' else 1) params['hidden_size'] = ih_shape[0] // multiplier params['input_size'] = ih_shape[1] return params def infer_shapes(layer_type: str, shapes: Dict[str, List[int]], params: Dict[str, Any]) -> Tuple[Optional[List[int]], Optional[List[int]]]: """Infer input/output shapes from layer parameters.""" input_shape = None output_shape = None if layer_type == 'Linear': if 'in_features' in params: input_shape = [-1, params['in_features']] if 'out_features' in params: output_shape = [-1, params['out_features']] elif layer_type in ('Conv2d',): if 'in_channels' in params: input_shape = [-1, params['in_channels'], -1, -1] if 'out_channels' in params: output_shape = [-1, params['out_channels'], -1, -1] elif layer_type in ('Conv1d',): if 'in_channels' in params: input_shape = [-1, params['in_channels'], -1] if 'out_channels' in params: output_shape = [-1, params['out_channels'], -1] elif layer_type in ('BatchNorm2d',): if 'num_features' in params: input_shape = [-1, params['num_features'], -1, -1] output_shape = [-1, params['num_features'], -1, -1] elif layer_type == 'Embedding': if 'embedding_dim' in params: output_shape = [-1, -1, params['embedding_dim']] elif layer_type in ('GRU', 'LSTM', 'RNN'): # For recurrent layers: input is (batch, seq_len, input_size) # output is (batch, seq_len, hidden_size * num_directions) if 'input_size' in params: input_shape = [-1, -1, params['input_size']] if 'hidden_size' in params: num_directions = 2 if params.get('bidirectional', False) else 1 output_shape = [-1, -1, params['hidden_size'] * num_directions] return input_shape, output_shape def architecture_to_dict(arch: ModelArchitecture) -> Dict[str, Any]: """Convert ModelArchitecture to JSON-serializable dict.""" return { 'name': arch.name, 'framework': arch.framework, 'totalParameters': arch.total_parameters, 'trainableParameters': arch.trainable_parameters, 'inputShape': arch.input_shape, 'outputShape': arch.output_shape, 'layers': [ { 'id': layer.id, 'name': layer.name, 'type': layer.type, 'category': layer.category, 'inputShape': layer.input_shape, 'outputShape': layer.output_shape, 'params': layer.params, 'numParameters': layer.num_parameters, 'trainable': layer.trainable } for layer in arch.layers ], 'connections': [ { 'source': conn.source, 'target': conn.target, 'tensorShape': conn.tensor_shape } for conn in arch.connections ] }