quantmllibrary 0.1.0__py3-none-any.whl

quantml/optim/rmsprop.py

@@ -0,0 +1,203 @@
+"""
+RMSProp optimizer implementation.
+
+RMSProp (Root Mean Square Propagation) is an adaptive learning rate optimizer
+that maintains a moving average of squared gradients.
+"""
+
+from typing import List, Optional, Dict, Any
+from quantml.tensor import Tensor
+from quantml import ops
+
+# Try to import NumPy
+try:
+    import numpy as np
+    HAS_NUMPY = True
+except ImportError:
+    HAS_NUMPY = False
+    np = None
+
+
+class RMSProp:
+    """
+    RMSProp optimizer.
+    
+    RMSProp maintains a moving average of squared gradients and divides
+    the gradient by the root of this average.
+    
+    Attributes:
+        lr: Learning rate
+        alpha: Smoothing constant (decay factor)
+        eps: Small value for numerical stability
+        weight_decay: Weight decay coefficient
+        momentum: Momentum factor (0 = no momentum)
+        squared_avg: Moving average of squared gradients
+    
+    Examples:
+        >>> optimizer = RMSProp(lr=0.001, alpha=0.99)
+        >>> for param in model.parameters():
+        >>>     optimizer.step(param)
+    """
+    
+    def __init__(
+        self,
+        params: Optional[List[Tensor]] = None,
+        lr: float = 0.01,
+        alpha: float = 0.99,
+        eps: float = 1e-8,
+        weight_decay: float = 0.0,
+        momentum: float = 0.0
+    ):
+        """
+        Initialize RMSProp optimizer.
+        
+        Args:
+            params: Optional list of parameters to optimize
+            lr: Learning rate
+            alpha: Smoothing constant (decay factor for squared gradient average)
+            eps: Small value to prevent division by zero
+            weight_decay: Weight decay (L2 regularization) coefficient
+            momentum: Momentum factor (0.0 to disable)
+        """
+        self.params = params if params is not None else []
+        self.lr = lr
+        self.alpha = alpha
+        self.eps = eps
+        self.weight_decay = weight_decay
+        self.momentum = momentum
+        self.squared_avg: Dict[int, Any] = {}  # Moving average of squared gradients
+        self.momentum_buffer: Dict[int, Any] = {}  # Momentum buffer
+    
+    def step(self, param: Optional[Tensor] = None):
+        """
+        Perform a single optimization step.
+        
+        Args:
+            param: Optional single parameter to update
+        """
+        if param is not None:
+            self._update_param(param)
+        else:
+            for p in self.params:
+                self._update_param(p)
+    
+    def _update_param(self, param: Tensor):
+        """Update a single parameter using RMSProp algorithm."""
+        if not param.requires_grad:
+            return
+        
+        if param.grad is None:
+            return
+        
+        param_id = id(param)
+        
+        if HAS_NUMPY:
+            try:
+                # Get gradient and parameter as NumPy arrays
+                grad = param.grad
+                if isinstance(grad, np.ndarray):
+                    grad_arr = grad
+                else:
+                    grad_arr = np.array(grad, dtype=np.float64)
+                
+                param_arr = param.numpy if param.numpy is not None else np.array(param.data, dtype=np.float64)
+                
+                # Apply weight decay
+                if self.weight_decay > 0:
+                    grad_arr = grad_arr + self.weight_decay * param_arr
+                
+                # Initialize squared average if needed
+                if param_id not in self.squared_avg:
+                    self.squared_avg[param_id] = np.zeros_like(param_arr, dtype=np.float64)
+                
+                # Update squared average: avg = alpha * avg + (1 - alpha) * grad^2
+                sq_avg = self.squared_avg[param_id]
+                sq_avg[:] = self.alpha * sq_avg + (1.0 - self.alpha) * (grad_arr ** 2)
+                
+                # Compute update: update = grad / (sqrt(avg) + eps)
+                update = grad_arr / (np.sqrt(sq_avg) + self.eps)
+                
+                # Apply momentum if enabled
+                if self.momentum > 0:
+                    if param_id not in self.momentum_buffer:
+                        self.momentum_buffer[param_id] = np.zeros_like(param_arr, dtype=np.float64)
+                    
+                    buf = self.momentum_buffer[param_id]
+                    buf[:] = self.momentum * buf + update
+                    update = buf
+                
+                # Update parameter: param = param - lr * update
+                param_update = self.lr * update
+                new_param_arr = param_arr - param_update
+                param.data = new_param_arr
+                
+            except (ValueError, TypeError, AttributeError):
+                self._update_param_fallback(param)
+        else:
+            self._update_param_fallback(param)
+    
+    def _update_param_fallback(self, param: Tensor):
+        """Fallback update using Tensor operations."""
+        if param.grad is None:
+            return
+        
+        param_id = id(param)
+        
+        # Initialize squared average if needed
+        if param_id not in self.squared_avg:
+            if isinstance(param.data[0], list):
+                self.squared_avg[param_id] = [[0.0] * len(row) for row in param.data]
+            else:
+                self.squared_avg[param_id] = [0.0] * len(param.data)
+        
+        grad = param.grad
+        if self.weight_decay > 0:
+            grad = ops.add(grad, ops.mul(param, self.weight_decay))
+        
+        # Update squared average
+        grad_sq = ops.mul(grad, grad)
+        sq_avg = Tensor(self.squared_avg[param_id])
+        new_sq_avg = ops.add(
+            ops.mul(sq_avg, self.alpha),
+            ops.mul(grad_sq, 1.0 - self.alpha)
+        )
+        self.squared_avg[param_id] = new_sq_avg.data
+        
+        # Compute update
+        sq_avg_sqrt = ops.pow(ops.add(new_sq_avg, self.eps), 0.5)
+        update = ops.div(grad, sq_avg_sqrt)
+        
+        # Apply momentum
+        if self.momentum > 0:
+            if param_id not in self.momentum_buffer:
+                if isinstance(param.data[0], list):
+                    self.momentum_buffer[param_id] = [[0.0] * len(row) for row in param.data]
+                else:
+                    self.momentum_buffer[param_id] = [0.0] * len(param.data)
+            
+            buf = Tensor(self.momentum_buffer[param_id])
+            new_buf = ops.add(ops.mul(buf, self.momentum), update)
+            self.momentum_buffer[param_id] = new_buf.data
+            update = new_buf
+        
+        # Update parameter
+        param_update = ops.mul(update, self.lr)
+        if param.requires_grad:
+            param_detached = param.detach()
+            param_detached.sub_(param_update)
+            param.data = param_detached.data
+        else:
+            param.sub_(param_update)
+    
+    def zero_grad(self, param: Optional[Tensor] = None):
+        """Clear gradients."""
+        if param is not None:
+            param.zero_grad()
+        else:
+            for p in self.params:
+                p.zero_grad()
+    
+    def add_param_group(self, params: List[Tensor]):
+        """Add a parameter group to optimize."""
+        self.params.extend(params)
+

quantml/optim/schedulers.py

@@ -0,0 +1,286 @@
+"""
+Learning rate schedulers for QuantML.
+
+Provides various learning rate scheduling strategies for training optimization.
+"""
+
+from typing import List, Optional, Any
+from abc import ABC, abstractmethod
+import math
+
+
+class LRScheduler(ABC):
+    """
+    Base class for learning rate schedulers.
+    
+    All schedulers should inherit from this class and implement
+    the get_lr() and step() methods.
+    """
+    
+    def __init__(self, optimizer: Any, last_epoch: int = -1):
+        """
+        Initialize scheduler.
+        
+        Args:
+            optimizer: The optimizer to schedule
+            last_epoch: The index of the last epoch
+        """
+        self.optimizer = optimizer
+        self.last_epoch = last_epoch
+        self.base_lrs = [group.get('lr', optimizer.lr) if isinstance(group, dict) else optimizer.lr 
+                        for group in getattr(optimizer, 'param_groups', [optimizer])]
+    
+    @abstractmethod
+    def get_lr(self) -> List[float]:
+        """Compute learning rate for current epoch."""
+        pass
+    
+    def step(self, epoch: Optional[int] = None):
+        """
+        Step the scheduler.
+        
+        Args:
+            epoch: Current epoch (if None, uses last_epoch + 1)
+        """
+        if epoch is None:
+            self.last_epoch += 1
+        else:
+            self.last_epoch = epoch
+        
+        lrs = self.get_lr()
+        # Update optimizer learning rates
+        if hasattr(self.optimizer, 'param_groups'):
+            for param_group, lr in zip(self.optimizer.param_groups, lrs):
+                param_group['lr'] = lr
+        else:
+            self.optimizer.lr = lrs[0] if lrs else self.optimizer.lr
+
+
+class StepLR(LRScheduler):
+    """Step learning rate scheduler - decays LR by gamma every step_size epochs."""
+    
+    def __init__(self, optimizer: Any, step_size: int, gamma: float = 0.1, last_epoch: int = -1):
+        """
+        Initialize StepLR scheduler.
+        
+        Args:
+            optimizer: The optimizer
+            step_size: Period of learning rate decay
+            gamma: Multiplicative factor for decay
+            last_epoch: The index of the last epoch
+        """
+        self.step_size = step_size
+        self.gamma = gamma
+        super().__init__(optimizer, last_epoch)
+    
+    def get_lr(self) -> List[float]:
+        """Compute learning rate."""
+        return [base_lr * (self.gamma ** (self.last_epoch // self.step_size)) 
+                for base_lr in self.base_lrs]
+
+
+class CosineAnnealingLR(LRScheduler):
+    """Cosine annealing learning rate scheduler."""
+    
+    def __init__(self, optimizer: Any, T_max: int, eta_min: float = 0.0, last_epoch: int = -1):
+        """
+        Initialize CosineAnnealingLR scheduler.
+        
+        Args:
+            optimizer: The optimizer
+            T_max: Maximum number of iterations
+            eta_min: Minimum learning rate
+            last_epoch: The index of the last epoch
+        """
+        self.T_max = T_max
+        self.eta_min = eta_min
+        super().__init__(optimizer, last_epoch)
+    
+    def get_lr(self) -> List[float]:
+        """Compute learning rate."""
+        return [self.eta_min + (base_lr - self.eta_min) * 
+                (1 + math.cos(math.pi * self.last_epoch / self.T_max)) / 2
+                for base_lr in self.base_lrs]
+
+
+class WarmupLR(LRScheduler):
+    """Warmup learning rate scheduler."""
+    
+    def __init__(self, optimizer: Any, warmup_steps: int, warmup_type: str = 'linear', last_epoch: int = -1):
+        """
+        Initialize WarmupLR scheduler.
+        
+        Args:
+            optimizer: The optimizer
+            warmup_steps: Number of warmup steps
+            warmup_type: Type of warmup ('linear' or 'cosine')
+            last_epoch: The index of the last epoch
+        """
+        self.warmup_steps = warmup_steps
+        self.warmup_type = warmup_type
+        super().__init__(optimizer, last_epoch)
+    
+    def get_lr(self) -> List[float]:
+        """Compute learning rate."""
+        if self.last_epoch < self.warmup_steps:
+            if self.warmup_type == 'linear':
+                factor = (self.last_epoch + 1) / self.warmup_steps
+            else:  # cosine
+                factor = (1 + math.cos(math.pi * (1 - (self.last_epoch + 1) / self.warmup_steps))) / 2
+            return [base_lr * factor for base_lr in self.base_lrs]
+        return self.base_lrs
+
+
+class ReduceLROnPlateau:
+    """Reduce learning rate when a metric has stopped improving."""
+    
+    def __init__(self, optimizer: Any, mode: str = 'min', factor: float = 0.1, 
+                 patience: int = 10, threshold: float = 1e-4, min_lr: float = 0.0):
+        """
+        Initialize ReduceLROnPlateau scheduler.
+        
+        Args:
+            optimizer: The optimizer
+            mode: 'min' or 'max' - whether to reduce LR when metric stops decreasing/increasing
+            factor: Factor to multiply LR by
+            patience: Number of epochs with no improvement before reducing LR
+            threshold: Threshold for measuring improvement
+            min_lr: Minimum learning rate
+        """
+        self.optimizer = optimizer
+        self.mode = mode
+        self.factor = factor
+        self.patience = patience
+        self.threshold = threshold
+        self.min_lr = min_lr
+        self.best = None
+        self.num_bad_epochs = 0
+        self.base_lrs = [optimizer.lr]
+    
+    def step(self, metrics: float):
+        """
+        Step the scheduler based on metrics.
+        
+        Args:
+            metrics: Current metric value
+        """
+        if self.best is None:
+            self.best = metrics
+        else:
+            if self.mode == 'min':
+                is_better = metrics < self.best - self.threshold
+            else:
+                is_better = metrics > self.best + self.threshold
+            
+            if is_better:
+                self.best = metrics
+                self.num_bad_epochs = 0
+            else:
+                self.num_bad_epochs += 1
+            
+            if self.num_bad_epochs >= self.patience:
+                self._reduce_lr()
+                self.num_bad_epochs = 0
+    
+    def _reduce_lr(self):
+        """Reduce learning rate."""
+        new_lr = max(self.optimizer.lr * self.factor, self.min_lr)
+        self.optimizer.lr = new_lr
+
+
+class CyclicLR(LRScheduler):
+    """Cyclic learning rate scheduler."""
+    
+    def __init__(self, optimizer: Any, base_lr: float, max_lr: float, 
+                 step_size_up: int = 2000, step_size_down: Optional[int] = None,
+                 mode: str = 'triangular', gamma: float = 1.0, last_epoch: int = -1):
+        """
+        Initialize CyclicLR scheduler.
+        
+        Args:
+            optimizer: The optimizer
+            base_lr: Lower bound of learning rate
+            max_lr: Upper bound of learning rate
+            step_size_up: Number of steps to increase LR
+            step_size_down: Number of steps to decrease LR (if None, equals step_size_up)
+            mode: 'triangular', 'triangular2', or 'exp_range'
+            gamma: Scaling factor for 'exp_range' mode
+            last_epoch: The index of the last epoch
+        """
+        self.base_lr = base_lr
+        self.max_lr = max_lr
+        self.step_size_up = step_size_up
+        self.step_size_down = step_size_down if step_size_down is not None else step_size_up
+        self.mode = mode
+        self.gamma = gamma
+        self.step_size = self.step_size_up + self.step_size_down
+        super().__init__(optimizer, last_epoch)
+    
+    def get_lr(self) -> List[float]:
+        """Compute learning rate."""
+        cycle = math.floor(1 + self.last_epoch / self.step_size)
+        x = 1 + self.last_epoch / self.step_size - cycle
+        
+        if x <= self.step_size_up / self.step_size:
+            scale = x * (self.step_size / self.step_size_up)
+        else:
+            scale = (self.step_size - x) * (self.step_size / self.step_size_down)
+        
+        if self.mode == 'triangular':
+            lr = self.base_lr + (self.max_lr - self.base_lr) * max(0, scale)
+        elif self.mode == 'triangular2':
+            lr = self.base_lr + (self.max_lr - self.base_lr) * max(0, scale) / (2 ** (cycle - 1))
+        else:  # exp_range
+            lr = self.base_lr + (self.max_lr - self.base_lr) * max(0, scale) * (self.gamma ** self.last_epoch)
+        
+        return [lr for _ in self.base_lrs]
+
+
+class OneCycleLR(LRScheduler):
+    """One cycle learning rate scheduler."""
+    
+    def __init__(self, optimizer: Any, max_lr: float, total_steps: int,
+                 pct_start: float = 0.3, anneal_strategy: str = 'cos',
+                 div_factor: float = 25.0, final_div_factor: float = 10000.0, last_epoch: int = -1):
+        """
+        Initialize OneCycleLR scheduler.
+        
+        Args:
+            optimizer: The optimizer
+            max_lr: Maximum learning rate
+            total_steps: Total number of steps
+            pct_start: Percentage of steps for warmup
+            anneal_strategy: 'cos' or 'linear' annealing
+            div_factor: Initial LR = max_lr / div_factor
+            final_div_factor: Final LR = initial_lr / final_div_factor
+            last_epoch: The index of the last epoch
+        """
+        self.max_lr = max_lr
+        self.total_steps = total_steps
+        self.pct_start = pct_start
+        self.anneal_strategy = anneal_strategy
+        self.div_factor = div_factor
+        self.final_div_factor = final_div_factor
+        self.initial_lr = max_lr / div_factor
+        self.final_lr = self.initial_lr / final_div_factor
+        super().__init__(optimizer, last_epoch)
+    
+    def get_lr(self) -> List[float]:
+        """Compute learning rate."""
+        if self.last_epoch < self.total_steps * self.pct_start:
+            # Warmup phase
+            pct = self.last_epoch / (self.total_steps * self.pct_start)
+            if self.anneal_strategy == 'cos':
+                lr = self.initial_lr + (self.max_lr - self.initial_lr) * (1 + math.cos(math.pi * (1 - pct))) / 2
+            else:
+                lr = self.initial_lr + (self.max_lr - self.initial_lr) * pct
+        else:
+            # Annealing phase
+            pct = (self.last_epoch - self.total_steps * self.pct_start) / (self.total_steps * (1 - self.pct_start))
+            if self.anneal_strategy == 'cos':
+                lr = self.final_lr + (self.max_lr - self.final_lr) * (1 + math.cos(math.pi * pct)) / 2
+            else:
+                lr = self.max_lr - (self.max_lr - self.final_lr) * pct
+        
+        return [lr for _ in self.base_lrs]
+

quantml/optim/sgd.py

@@ -0,0 +1,181 @@
+"""
+Stochastic Gradient Descent (SGD) optimizer.
+
+This module provides the SGD optimizer with optional momentum and weight decay.
+"""
+
+from typing import List, Optional, Dict, Any
+from quantml.tensor import Tensor
+from quantml import ops
+
+# Try to import NumPy
+try:
+    import numpy as np
+    HAS_NUMPY = True
+except ImportError:
+    HAS_NUMPY = False
+    np = None
+
+
+class SGD:
+    """
+    Stochastic Gradient Descent optimizer.
+    
+    Updates parameters using: param = param - lr * (grad + weight_decay * param)
+    With momentum: v = momentum * v + grad, param = param - lr * v
+    
+    Attributes:
+        lr: Learning rate
+        momentum: Momentum factor (0 = no momentum)
+        weight_decay: Weight decay (L2 regularization) factor
+        velocity: Momentum velocity for each parameter
+    
+    Examples:
+        >>> optimizer = SGD(lr=0.01, momentum=0.9)
+        >>> for param in model.parameters():
+        >>>     optimizer.step(param)
+    """
+    
+    def __init__(
+        self,
+        params: Optional[List[Tensor]] = None,
+        lr: float = 0.01,
+        momentum: float = 0.0,
+        weight_decay: float = 0.0
+    ):
+        """
+        Initialize SGD optimizer.
+        
+        Args:
+            params: Optional list of parameters to optimize
+            lr: Learning rate
+            momentum: Momentum factor (0.0 to disable)
+            weight_decay: Weight decay coefficient
+        """
+        self.params = params if params is not None else []
+        self.lr = lr
+        self.momentum = momentum
+        self.weight_decay = weight_decay
+        self.velocity: Dict[int, Any] = {}  # Store velocity as NumPy arrays
+    
+    def step(self, param: Optional[Tensor] = None):
+        """
+        Perform a single optimization step.
+        
+        If param is provided, updates that parameter.
+        Otherwise, updates all parameters in self.params.
+        
+        Args:
+            param: Optional single parameter to update
+        """
+        if param is not None:
+            self._update_param(param)
+        else:
+            for p in self.params:
+                self._update_param(p)
+    
+    def _update_param(self, param: Tensor):
+        """Update a single parameter using direct NumPy operations."""
+        if not param.requires_grad:
+            return
+        
+        if param.grad is None:
+            return
+        
+        param_id = id(param)
+        
+        # Get gradient as NumPy array if possible
+        if HAS_NUMPY:
+            try:
+                # Get gradient as NumPy array
+                grad = param.grad
+                if isinstance(grad, np.ndarray):
+                    grad_arr = grad
+                else:
+                    grad_arr = np.array(grad, dtype=np.float64)
+                
+                # Get parameter as NumPy array
+                param_arr = param.numpy if param.numpy is not None else np.array(param.data, dtype=np.float64)
+                
+                # Apply weight decay
+                if self.weight_decay > 0:
+                    grad_arr = grad_arr + self.weight_decay * param_arr
+                
+                # Update velocity if momentum is used
+                if self.momentum > 0:
+                    if param_id not in self.velocity:
+                        # Initialize velocity to zero
+                        self.velocity[param_id] = np.zeros_like(param_arr, dtype=np.float64)
+                    
+                    # v = momentum * v + grad
+                    vel = self.velocity[param_id]
+                    vel[:] = self.momentum * vel + grad_arr
+                    update = vel
+                else:
+                    update = grad_arr
+                
+                # Compute parameter update: param = param - lr * update
+                param_update = self.lr * update
+                new_param_arr = param_arr - param_update
+                
+                # Update parameter data directly
+                param.data = new_param_arr
+                
+            except (ValueError, TypeError, AttributeError):
+                # Fallback to Tensor operations
+                self._update_param_fallback(param)
+        else:
+            # Fallback to Tensor operations
+            self._update_param_fallback(param)
+    
+    def _update_param_fallback(self, param: Tensor):
+        """Fallback update using Tensor operations."""
+        if param.grad is None:
+            return
+        
+        param_id = id(param)
+        if param_id not in self.velocity:
+            if isinstance(param.data[0], list):
+                self.velocity[param_id] = [[0.0] * len(row) for row in param.data]
+            else:
+                self.velocity[param_id] = [0.0] * len(param.data)
+        
+        grad = param.grad
+        if self.weight_decay > 0:
+            grad = ops.add(grad, ops.mul(param, self.weight_decay))
+        
+        if self.momentum > 0:
+            vel = self.velocity[param_id]
+            vel_tensor = Tensor(vel)
+            new_vel = ops.add(ops.mul(vel_tensor, self.momentum), grad)
+            self.velocity[param_id] = new_vel.data
+            update = new_vel
+        else:
+            update = grad
+        
+        if param.requires_grad:
+            param_detached = param.detach()
+            param_update = ops.mul(update, self.lr)
+            param_detached.sub_(param_update)
+            param.data = param_detached.data
+        else:
+            param_update = ops.mul(update, self.lr)
+            param.sub_(param_update)
+    
+    def zero_grad(self, param: Optional[Tensor] = None):
+        """
+        Clear gradients.
+        
+        Args:
+            param: Optional single parameter, otherwise clears all
+        """
+        if param is not None:
+            param.zero_grad()
+        else:
+            for p in self.params:
+                p.zero_grad()
+    
+    def add_param_group(self, params: List[Tensor]):
+        """Add a parameter group to optimize."""
+        self.params.extend(params)
+