optimaxx 0.1.0__py3-none-any.whl

optimax/__init__.py

@@ -0,0 +1,21 @@
+"""Optimax: An intelligent AI model optimization framework for PyTorch."""
+
+from optimax.analyzer import analyze
+from optimax.optimizer import optimize, OptimizationResult
+from optimax.benchmark import benchmark
+from optimax.hardware import detect_hardware
+from optimax.report import Report, generate_report
+from optimax.config import load_config
+from optimax.version import __version__
+
+__all__ = [
+    "analyze",
+    "optimize",
+    "benchmark",
+    "detect_hardware",
+    "generate_report",
+    "load_config",
+    "OptimizationResult",
+    "Report",
+    "__version__",
+]

optimax/analyzer/__init__.py

@@ -0,0 +1,24 @@
+"""Analyzer module for Optimax."""
+
+from optimax.analyzer.analyzer import analyze, ModelAnalyzer
+from optimax.analyzer.models import ModelProfile, LayerInfo, TensorProfile
+from optimax.analyzer.recommendation import (
+    generate_recommendations,
+    Recommendation,
+    RecommendationEngine,
+    RiskLevel,
+    Compatibility,
+)
+
+__all__ = [
+    "analyze",
+    "ModelAnalyzer",
+    "ModelProfile",
+    "LayerInfo",
+    "TensorProfile",
+    "generate_recommendations",
+    "Recommendation",
+    "RecommendationEngine",
+    "RiskLevel",
+    "Compatibility",
+]

optimax/analyzer/analyzer.py

@@ -0,0 +1,255 @@
+"""Model analysis for Optimax."""
+
+from typing import Any, Optional
+
+import torch
+
+from optimax.analyzer.models import LayerInfo, ModelProfile, TensorProfile
+from optimax.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+class ModelAnalyzer:
+    """Analyzes PyTorch models and collects comprehensive statistics."""
+
+    def __init__(self, model: torch.nn.Module) -> None:
+        """Initialize the analyzer with a model.
+
+        Args:
+            model: PyTorch model to analyze.
+        """
+        self.model = model
+        self._hooks: list[Any] = []
+        self._activation_shapes: dict[str, list[tuple[int, ...]]] = {}
+
+    def _register_forward_hooks(self) -> None:
+        """Register forward hooks to capture activation shapes."""
+        self._activation_shapes.clear()
+
+        def hook(name: str) -> Any:
+            def fn(_module: torch.nn.Module, input: Any, output: Any) -> None:
+                shapes: list[tuple[int, ...]] = []
+                if isinstance(output, torch.Tensor):
+                    shapes.append(tuple(output.shape))
+                elif isinstance(output, (list, tuple)):
+                    for o in output:
+                        if isinstance(o, torch.Tensor):
+                            shapes.append(tuple(o.shape))
+                self._activation_shapes[name] = shapes
+
+            return fn
+
+        for name, module in self.model.named_modules():
+            if name:
+                h = module.register_forward_hook(hook(name))
+                self._hooks.append(h)
+
+    def _remove_hooks(self) -> None:
+        """Remove all registered forward hooks."""
+        for h in self._hooks:
+            h.remove()
+        self._hooks.clear()
+
+    @staticmethod
+    def _count_flops_linear(module: torch.nn.Linear, input_shape: tuple[int, ...]) -> int:
+        """Estimate FLOPs for a linear layer.
+
+        Args:
+            module: Linear module.
+            input_shape: Input tensor shape.
+
+        Returns:
+            Estimated FLOPs.
+        """
+        # FLOPs = 2 * in_features * out_features * batch_size (assuming matmul)
+        in_features = module.in_features
+        out_features = module.out_features
+        batch_size = int(torch.prod(torch.tensor(input_shape[:-1])))
+        return 2 * in_features * out_features * batch_size
+
+    @staticmethod
+    def _count_flops_conv(module: torch.nn.modules.conv._ConvNd, input_shape: tuple[int, ...]) -> int:
+        """Estimate FLOPs for a convolution layer.
+
+        Args:
+            module: Conv module.
+            input_shape: Input tensor shape (N, C_in, ...).
+
+        Returns:
+            Estimated FLOPs.
+        """
+        # FLOPs = 2 * output_elements * kernel_elements * C_in / groups
+        output_shape = module(output_shape=input_shape)  # type: ignore[arg-type]
+        # Simpler estimation: use formula
+        kernel_dims = list(module.kernel_size)
+        kernel_elements = int(torch.prod(torch.tensor(kernel_dims)))
+        in_channels = module.in_channels
+        out_channels = module.out_channels
+        groups = module.groups
+        # Output spatial size
+        if len(input_shape) == 4:  # 2D conv
+            batch_size, _, h, w = input_shape
+            out_h = (h + 2 * module.padding[0] - module.dilation[0] * (kernel_dims[0] - 1) - 1) // module.stride[0] + 1
+            out_w = (w + 2 * module.padding[1] - module.dilation[1] * (kernel_dims[1] - 1) - 1) // module.stride[1] + 1
+            output_elements = batch_size * out_h * out_w
+        elif len(input_shape) == 5:  # 3D conv
+            batch_size, _, d, h, w = input_shape
+            out_d = (d + 2 * module.padding[0] - module.dilation[0] * (kernel_dims[0] - 1) - 1) // module.stride[0] + 1
+            out_h = (h + 2 * module.padding[1] - module.dilation[1] * (kernel_dims[1] - 1) - 1) // module.stride[1] + 1
+            out_w = (w + 2 * module.padding[2] - module.dilation[2] * (kernel_dims[2] - 1) - 1) // module.stride[2] + 1
+            output_elements = batch_size * out_d * out_h * out_w
+        elif len(input_shape) == 3:  # 1D conv
+            batch_size, _, l = input_shape
+            out_l = (l + 2 * module.padding[0] - module.dilation[0] * (kernel_dims[0] - 1) - 1) // module.stride[0] + 1
+            output_elements = batch_size * out_l
+        else:
+            return 0
+        return 2 * output_elements * kernel_elements * in_channels * out_channels // groups
+
+    def _estimate_module_flops(
+        self, module: torch.nn.Module, input_shapes: list[tuple[int, ...]]
+    ) -> Optional[int]:
+        """Estimate FLOPs for a single module.
+
+        Args:
+            module: A PyTorch module.
+            input_shapes: List of input tensor shapes.
+
+        Returns:
+            Estimated FLOPs or None if not supported.
+        """
+        if not input_shapes:
+            return None
+        input_shape = input_shapes[0]
+        if isinstance(module, torch.nn.Linear):
+            return self._count_flops_linear(module, input_shape)
+        if isinstance(module, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d)):
+            return self._count_flops_conv(module, input_shape)
+        return None
+
+    def _analyze_module(self, name: str, module: torch.nn.Module) -> LayerInfo:
+        """Recursively analyze a module and its children.
+
+        Args:
+            name: Module name.
+            module: PyTorch module.
+
+        Returns:
+            LayerInfo with statistics.
+        """
+        params = sum(p.numel() for p in module.parameters())
+        trainable = sum(p.numel() for p in module.parameters() if p.requires_grad)
+        frozen = params - trainable
+        memory_mb = sum(p.numel() * p.element_size() for p in module.parameters()) / (1024 * 1024)
+        input_shapes = self._activation_shapes.get(name, [])
+        flops = self._estimate_module_flops(module, input_shapes)
+        children: list[LayerInfo] = []
+        for child_name, child_module in module.named_children():
+            full_name = f"{name}.{child_name}" if name else child_name
+            children.append(self._analyze_module(full_name, child_module))
+        return LayerInfo(
+            name=name,
+            type=module.__class__.__name__,
+            parameters=params,
+            trainable_parameters=trainable,
+            frozen_parameters=frozen,
+            input_shapes=[tuple(s) for s in input_shapes],
+            memory_mb=memory_mb,
+            flops=flops,
+            children=children,
+        )
+
+    def analyze(
+        self,
+        example_input: Optional[torch.Tensor] = None,
+    ) -> ModelProfile:
+        """Analyze the model and return a comprehensive profile.
+
+        Args:
+            example_input: Optional example input tensor for activation tracing.
+
+        Returns:
+            A ModelProfile instance.
+        """
+        total_params = sum(p.numel() for p in self.model.parameters())
+        trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
+        frozen_params = total_params - trainable_params
+        model_size = sum(p.numel() * p.element_size() for p in self.model.parameters()) / (1024 * 1024)
+        estimated_memory = model_size  # Base estimate
+
+        if example_input is not None:
+            self._register_forward_hooks()
+            try:
+                self.model(example_input)
+            except Exception as e:
+                logger.warning("Forward pass with example_input failed: %s", e)
+            self._remove_hooks()
+
+        layer_hierarchy = []
+        for name, module in self.model.named_children():
+            layer_hierarchy.append(self._analyze_module(name, module))
+
+        layer_count = sum(1 for _ in self.model.modules()) - 1  # exclude root
+
+        # Estimate total FLOPs from layer hierarchy
+        estimated_flops: Optional[int] = None
+        try:
+            total_flops = 0
+            def sum_flops(layers: list[LayerInfo]) -> None:
+                nonlocal total_flops
+                for layer in layers:
+                    if layer.flops is not None:
+                        total_flops += layer.flops
+                    sum_flops(layer.children)
+            sum_flops(layer_hierarchy)
+            estimated_flops = total_flops
+        except Exception:
+            pass
+
+        tensor_profiles: list[TensorProfile] = []
+        for shapes in self._activation_shapes.values():
+            for shape in shapes:
+                elements = int(torch.prod(torch.tensor(shape)))
+                memory_bytes = elements * 4  # assume float32
+                tensor_profiles.append(
+                    TensorProfile(
+                        shape=shape,
+                        dtype="float32",
+                        elements=elements,
+                        memory_bytes=memory_bytes,
+                    )
+                )
+
+        # Add activation memory to estimated memory
+        activation_memory = sum(t.memory_bytes for t in tensor_profiles) / (1024 * 1024)
+        estimated_memory += activation_memory
+
+        return ModelProfile(
+            total_parameters=total_params,
+            trainable_parameters=trainable_params,
+            frozen_parameters=frozen_params,
+            model_size_mb=model_size,
+            estimated_memory_mb=estimated_memory,
+            estimated_flops=estimated_flops,
+            layer_hierarchy=layer_hierarchy,
+            layer_count=layer_count,
+            tensor_profiles=tensor_profiles,
+        )
+
+
+def analyze(
+    model: torch.nn.Module,
+    example_input: Optional[torch.Tensor] = None,
+) -> ModelProfile:
+    """Convenience function to analyze a model.
+
+    Args:
+        model: PyTorch model to analyze.
+        example_input: Optional example input for activation tracing.
+
+    Returns:
+        A ModelProfile instance.
+    """
+    analyzer = ModelAnalyzer(model)
+    return analyzer.analyze(example_input)

optimax/analyzer/models.py

@@ -0,0 +1,85 @@
+"""Data models for the analyzer."""
+
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+import torch
+
+
+@dataclass
+class LayerInfo:
+    """Information about a single layer/module."""
+
+    name: str
+    type: str
+    parameters: int
+    trainable_parameters: int
+    frozen_parameters: int
+    input_shapes: list[Optional[tuple[int, ...]]] = field(default_factory=list)
+    output_shapes: list[Optional[tuple[int, ...]]] = field(default_factory=list)
+    memory_mb: float = 0.0
+    flops: Optional[int] = None
+    children: list["LayerInfo"] = field(default_factory=list)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Return dictionary representation."""
+        return {
+            "name": self.name,
+            "type": self.type,
+            "parameters": self.parameters,
+            "trainable_parameters": self.trainable_parameters,
+            "frozen_parameters": self.frozen_parameters,
+            "input_shapes": [list(s) if s else None for s in self.input_shapes],
+            "output_shapes": [list(s) if s else None for s in self.output_shapes],
+            "memory_mb": self.memory_mb,
+            "flops": self.flops,
+            "children": [c.to_dict() for c in self.children],
+        }
+
+
+@dataclass
+class TensorProfile:
+    """Profile of tensor activations."""
+
+    shape: tuple[int, ...]
+    dtype: str
+    elements: int
+    memory_bytes: int
+
+    def to_dict(self) -> dict[str, Any]:
+        """Return dictionary representation."""
+        return {
+            "shape": list(self.shape),
+            "dtype": self.dtype,
+            "elements": self.elements,
+            "memory_bytes": self.memory_bytes,
+        }
+
+
+@dataclass
+class ModelProfile:
+    """Complete profile of a PyTorch model."""
+
+    total_parameters: int
+    trainable_parameters: int
+    frozen_parameters: int
+    model_size_mb: float
+    estimated_memory_mb: float
+    estimated_flops: Optional[int] = None
+    layer_hierarchy: list[LayerInfo] = field(default_factory=list)
+    layer_count: int = 0
+    tensor_profiles: list[TensorProfile] = field(default_factory=list)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Return dictionary representation."""
+        return {
+            "total_parameters": self.total_parameters,
+            "trainable_parameters": self.trainable_parameters,
+            "frozen_parameters": self.frozen_parameters,
+            "model_size_mb": self.model_size_mb,
+            "estimated_memory_mb": self.estimated_memory_mb,
+            "estimated_flops": self.estimated_flops,
+            "layer_hierarchy": [l.to_dict() for l in self.layer_hierarchy],
+            "layer_count": self.layer_count,
+            "tensor_profiles": [t.to_dict() for t in self.tensor_profiles],
+        }

optimax/analyzer/recommendation.py

@@ -0,0 +1,298 @@
+"""Recommendation engine for Optimax."""
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Optional
+
+import torch
+
+from optimax.analyzer.models import ModelProfile
+from optimax.hardware.models import HardwareProfile
+from optimax.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+class RiskLevel(str, Enum):
+    """Risk level for an optimization recommendation."""
+
+    LOW = "low"
+    MEDIUM = "medium"
+    HIGH = "high"
+
+
+class Compatibility(str, Enum):
+    """Compatibility assessment for a recommendation."""
+
+    FULL = "full"
+    PARTIAL = "partial"
+    LIMITED = "limited"
+
+
+@dataclass
+class Recommendation:
+    """A single optimization recommendation."""
+
+    name: str
+    description: str
+    reason: str
+    expected_speedup: Optional[float] = None
+    expected_memory_savings: Optional[float] = None
+    risk_level: RiskLevel = RiskLevel.LOW
+    compatibility: Compatibility = Compatibility.FULL
+    pass_name: Optional[str] = None
+    confidence: float = 0.5
+    supported_hardware: list[str] = field(default_factory=list)
+    required_dependencies: list[str] = field(default_factory=list)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Return dictionary representation."""
+        return {
+            "name": self.name,
+            "description": self.description,
+            "reason": self.reason,
+            "expected_speedup": self.expected_speedup,
+            "expected_memory_savings": self.expected_memory_savings,
+            "risk_level": self.risk_level.value,
+            "compatibility": self.compatibility.value,
+            "pass_name": self.pass_name,
+            "confidence": self.confidence,
+            "supported_hardware": self.supported_hardware,
+            "required_dependencies": self.required_dependencies,
+        }
+
+
+class RecommendationEngine:
+    """Generates optimization recommendations based on model and hardware analysis."""
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        profile: ModelProfile,
+        hardware: HardwareProfile,
+    ) -> None:
+        """Initialize the recommendation engine.
+
+        Args:
+            model: The PyTorch model.
+            profile: The model profile from analysis.
+            hardware: The detected hardware profile.
+        """
+        self.model = model
+        self.profile = profile
+        self.hardware = hardware
+        self._recommendations: list[Recommendation] = []
+
+    def _recommend_torch_compile(self) -> Optional[Recommendation]:
+        """Recommend torch.compile if compatible."""
+        if self.hardware.software.pytorch_version.startswith("2."):
+            reason = "PyTorch 2.x supports torch.compile() which can significantly reduce Python overhead and fuse kernels."
+            confidence = 0.8
+            if self.hardware.has_cuda:
+                confidence = 0.9
+            return Recommendation(
+                name="Enable torch.compile",
+                description="Use torch.compile() to optimize model execution.",
+                reason=reason,
+                expected_speedup=1.2,
+                expected_memory_savings=None,
+                risk_level=RiskLevel.LOW,
+                compatibility=Compatibility.FULL,
+                pass_name="torch_compile",
+                confidence=confidence,
+                supported_hardware=["cpu", "cuda", "mps"],
+                required_dependencies=["torch>=2.0"],
+            )
+        return None
+
+    def _recommend_mixed_precision(self) -> Optional[Recommendation]:
+        """Recommend mixed precision if GPU supports it."""
+        if self.hardware.has_cuda or self.hardware.has_metal:
+            reason = "Modern GPUs support Tensor Cores/Metal performance shaders that can accelerate fp16/bf16 operations significantly."
+            speedup = 1.5 if self.hardware.has_cuda else 1.2
+            memory = 0.5
+            confidence = 0.85 if self.hardware.has_cuda else 0.7
+            return Recommendation(
+                name="Enable Mixed Precision",
+                description="Use torch.autocast with fp16 or bf16.",
+                reason=reason,
+                expected_speedup=speedup,
+                expected_memory_savings=memory,
+                risk_level=RiskLevel.LOW,
+                compatibility=Compatibility.FULL,
+                pass_name="mixed_precision",
+                confidence=confidence,
+                supported_hardware=["cuda", "mps"],
+                required_dependencies=["torch>=2.0"],
+            )
+        return None
+
+    def _recommend_flash_attention(self) -> Optional[Recommendation]:
+        """Recommend FlashAttention if model has attention patterns."""
+        has_attention = any(
+            "attention" in layer.name.lower() or "attn" in layer.type.lower()
+            for layer in self.profile.layer_hierarchy
+        )
+        if has_attention and (self.hardware.has_cuda or self.hardware.has_metal):
+            reason = "FlashAttention reduces attention memory usage from O(N^2) to O(N) and is compute-bound on modern GPUs."
+            return Recommendation(
+                name="Enable FlashAttention",
+                description="Replace standard attention with FlashAttention kernels.",
+                reason=reason,
+                expected_speedup=1.3,
+                expected_memory_savings=0.6,
+                risk_level=RiskLevel.MEDIUM,
+                compatibility=Compatibility.PARTIAL,
+                pass_name="flash_attention",
+                confidence=0.7,
+                supported_hardware=["cuda", "mps"],
+                required_dependencies=["torch>=2.0"],
+            )
+        return None
+
+    def _recommend_quantization(self) -> Optional[Recommendation]:
+        """Recommend quantization for large models."""
+        if self.profile.total_parameters > 50_000_000:  # 50M params
+            reason = "Large models benefit from INT8/INT4 quantization for inference, reducing memory footprint and improving throughput."
+            return Recommendation(
+                name="Use Quantization",
+                description="Apply post-training quantization or QAT.",
+                reason=reason,
+                expected_speedup=1.5,
+                expected_memory_savings=0.75,
+                risk_level=RiskLevel.MEDIUM,
+                compatibility=Compatibility.PARTIAL,
+                pass_name="quantization",
+                confidence=0.6,
+                supported_hardware=["cpu", "cuda"],
+                required_dependencies=["torch"],
+            )
+        return None
+
+    def _recommend_operator_fusion(self) -> Optional[Recommendation]:
+        """Recommend operator fusion if there are many small ops."""
+        if self.profile.layer_count > 50:
+            reason = "Models with many small layers benefit from operator fusion, reducing kernel launch overhead and improving memory locality."
+            return Recommendation(
+                name="Fuse Operators",
+                description="Fuse adjacent compatible operations.",
+                reason=reason,
+                expected_speedup=1.1,
+                expected_memory_savings=0.1,
+                risk_level=RiskLevel.LOW,
+                compatibility=Compatibility.FULL,
+                pass_name="operator_fusion",
+                confidence=0.75,
+                supported_hardware=["cpu", "cuda", "mps"],
+                required_dependencies=["torch>=2.0"],
+            )
+        return None
+
+    def _recommend_channels_last(self) -> Optional[Recommendation]:
+        """Recommend channels-last memory layout for CNNs."""
+        has_conv = any(
+            "conv" in layer.type.lower()
+            for layer in self.profile.layer_hierarchy
+        )
+        if has_conv and (self.hardware.has_cuda or self.hardware.has_metal):
+            reason = "Channels-last (NHWC) memory layout can improve performance on GPUs with Tensor Cores for convolutional networks."
+            return Recommendation(
+                name="Use Channels-Last Layout",
+                description="Convert conv inputs to channels-last format.",
+                reason=reason,
+                expected_speedup=1.1,
+                expected_memory_savings=None,
+                risk_level=RiskLevel.LOW,
+                compatibility=Compatibility.PARTIAL,
+                pass_name="channels_last",
+                confidence=0.7,
+                supported_hardware=["cuda", "mps"],
+                required_dependencies=["torch"],
+            )
+        return None
+
+    def _recommend_gradient_checkpointing(self) -> Optional[Recommendation]:
+        """Recommend gradient checkpointing for large models with limited memory."""
+        memory_pressure = self.hardware.memory.percent_used > 80.0
+        large_model = self.profile.estimated_memory_mb > self.hardware.memory.available_mb * 0.5
+        if memory_pressure and large_model:
+            reason = "Gradient checkpointing trades compute for memory by recomputing activations during backward pass instead of storing them."
+            return Recommendation(
+                name="Use Gradient Checkpointing",
+                description="Enable torch.utils.checkpoint for memory-constrained training.",
+                reason=reason,
+                expected_speedup=None,
+                expected_memory_savings=0.5,
+                risk_level=RiskLevel.MEDIUM,
+                compatibility=Compatibility.PARTIAL,
+                pass_name="gradient_checkpointing",
+                confidence=0.8,
+                supported_hardware=["cpu", "cuda", "mps"],
+                required_dependencies=["torch"],
+            )
+        return None
+
+    def _recommend_reduce_batch_size(self) -> Optional[Recommendation]:
+        """Recommend reducing batch size if memory is constrained."""
+        if self.hardware.memory.percent_used > 90.0:
+            reason = "System memory is under high pressure. Reducing batch size is a straightforward way to decrease peak memory usage."
+            return Recommendation(
+                name="Reduce Batch Size",
+                description="Lower batch size to fit within available memory.",
+                reason=reason,
+                expected_speedup=None,
+                expected_memory_savings=0.3,
+                risk_level=RiskLevel.LOW,
+                compatibility=Compatibility.FULL,
+                pass_name=None,
+                confidence=0.9,
+                supported_hardware=["cpu", "cuda", "mps"],
+                required_dependencies=[],
+            )
+        return None
+
+    def generate(self) -> list[Recommendation]:
+        """Generate all applicable recommendations.
+
+        Returns:
+            A list of Recommendation instances sorted by confidence descending.
+        """
+        generators = [
+            self._recommend_torch_compile,
+            self._recommend_mixed_precision,
+            self._recommend_flash_attention,
+            self._recommend_quantization,
+            self._recommend_operator_fusion,
+            self._recommend_channels_last,
+            self._recommend_gradient_checkpointing,
+            self._recommend_reduce_batch_size,
+        ]
+        recommendations = []
+        for gen in generators:
+            try:
+                rec = gen()
+                if rec is not None:
+                    recommendations.append(rec)
+            except Exception as e:
+                logger.warning("Recommendation generator %s failed: %s", gen.__name__, e)
+        recommendations.sort(key=lambda r: r.confidence, reverse=True)
+        return recommendations
+
+
+def generate_recommendations(
+    model: torch.nn.Module,
+    profile: ModelProfile,
+    hardware: HardwareProfile,
+) -> list[Recommendation]:
+    """Convenience function to generate recommendations.
+
+    Args:
+        model: PyTorch model.
+        profile: Model profile.
+        hardware: Hardware profile.
+
+    Returns:
+        List of recommendations.
+    """
+    engine = RecommendationEngine(model=model, profile=profile, hardware=hardware)
+    return engine.generate()