Merge ee282be91f into 3265f2edfb

library/network_utils.py

@@ -0,0 +1,128 @@
+from contextlib import contextmanager
+import torch
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def maybe_sample_params(optimizer):
+    """
+    Returns parameter sampling context for IVON optimizers, otherwise returns no-op context.
+
+    pip install ivon-opt
+
+    Args:
+        optimizer: PyTorch optimizer instance.
+
+    Returns:
+        Context manager for parameter sampling if optimizer supports it, otherwise nullcontext().
+    """
+    from contextlib import nullcontext
+
+    return optimizer.sampled_params(train=True) if hasattr(optimizer, "sampled_params") else nullcontext()
+
+
+@contextmanager
+def maybe_pruned_save(model, optimizer, enable_pruning=False, pruning_ratio=0.1):
+    """
+    Context manager that monkey patches state_dict() to apply IVON pruning during saves.
+
+    Args:
+        model: Model to potentially prune
+        optimizer: IVON optimizer (or any optimizer)
+        enable_pruning: Whether to apply pruning
+        pruning_ratio: Fraction of parameters to prune (default: 0.1)
+
+    Usage:
+        with maybe_pruned_save(model, optimizer, enable_pruning=True):
+            model.save_weights(...)  # Saved state_dict will have pruned weights
+        # Model's state_dict is automatically restored after save
+    """
+    # Check if we should prune - more flexible detection of IVON-like optimizers
+    should_prune = enable_pruning and (
+        hasattr(optimizer, "sampled_params")
+    )
+
+    if not should_prune:
+        yield
+        return
+
+    param_variances = []
+
+    # Extract variances from IVON optimizer
+    offset = 0
+    for group in optimizer.param_groups:
+        # Get group-level values
+        ess = group["ess"]          # λ (lambda)
+        weight_decay = group["weight_decay"]  # δ (delta)
+        hess = group["hess"]        # hᵢ (Hessian diagonal)
+        
+        # Calculate variance: vᵢ = 1 / (λ × (hᵢ + δ))
+        group_variance = 1.0 / (ess * (hess + weight_decay))
+        
+        # Map back to individual parameters
+        param_offset = 0
+        for param in group["params"]:
+            if param is not None and param.requires_grad:
+                param_numel = param.numel()
+                param_slice = slice(param_offset, param_offset + param_numel)
+                
+                # Get variance for this parameter
+                param_var = group_variance[param_slice]
+                
+                # Store each element's variance with its location
+                flat_param_var = param_var.view(-1)
+                for i, var_val in enumerate(flat_param_var):
+                    param_variances.append((var_val.item(), param, i))
+                
+                param_offset += param_numel
+        
+        offset += group["numel"]
+    
+    if not param_variances:
+        yield
+        return
+    
+    param_variances.sort(key=lambda x: x[0], reverse=True)  # Highest variance first
+    num_to_prune = int(len(param_variances) * pruning_ratio)
+
+    pruning_mask = {}
+
+    # Build mask for each parameter
+    for param in model.parameters():
+        pruning_mask[id(param)] = torch.ones_like(param, dtype=torch.bool)
+
+    # Mark parameters to prune
+    for param in model.parameters():
+        mask = pruning_mask[id(param)]
+        num_to_prune = int(mask.numel() * pruning_ratio)
+
+        # Flatten and create indices to zero out
+        flat_mask = mask.view(-1)
+        prune_indices = torch.randperm(flat_mask.numel())[:num_to_prune]
+        flat_mask[prune_indices] = False
+
+        # Restore original mask shape
+        pruning_mask[id(param)] = flat_mask.view(mask.shape)
+
+    # Monkey patch state_dict
+    original_state_dict = model.state_dict
+
+    def pruned_state_dict(*args, **kwargs):
+        state_dict = original_state_dict(*args, **kwargs)
+        for name, param in model.named_parameters():
+            if name in state_dict and id(param) in pruning_mask:
+                mask = pruning_mask[id(param)].to(state_dict[name].device)
+                state_dict[name] = state_dict[name] * mask.float()
+        return state_dict
+
+    model.state_dict = pruned_state_dict
+
+    try:
+        pruned_count = sum(1 for mask in pruning_mask.values() for val in mask.flatten() if not val)
+        total_params = sum(mask.numel() for mask in pruning_mask.values())
+        logger.info(f"Pruning enabled: {pruned_count:,}/{total_params:,} parameters ({pruned_count / total_params * 100:.1f}%)")
+        yield
+    finally:
+        # Restore original state_dict
+        model.state_dict = original_state_dict