potnn 1.0.0__py3-none-any.whl

potnn/modules/depthwise.py

@@ -0,0 +1,216 @@
+"""PoT-quantized Depthwise Conv2d layer with Integer Simulation.
+
+v2: Added integer simulation for C-compatible QAT
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Union, Tuple
+
+from .base import PoTLayerBase
+from ..quantize.pot import quantize_to_pot_ste, quantize_to_pot, quantize_activation_ste
+from ..quantize.integer_sim import (
+    round_ste, floor_ste, clamp_ste,
+    quantize_to_int8_ste, quantize_to_uint8_ste,
+    requantize_ste
+)
+
+
+class PoTDepthwiseConv2d(PoTLayerBase):
+    """Power-of-Two quantized Depthwise Conv2d layer.
+
+    Depthwise convolution applies a single filter per input channel.
+    This is commonly used in MobileNet-style architectures as the first
+    part of depthwise separable convolution.
+    
+    Key properties:
+        - in_channels == out_channels == channels
+        - groups = channels (each channel processed independently)
+        - weight shape: [channels, 1, kH, kW]
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        kernel_size: Union[int, Tuple[int, int]] = 3,
+        stride: Union[int, Tuple[int, int]] = 1,
+        padding: Union[int, Tuple[int, int]] = 1,
+        dilation: Union[int, Tuple[int, int]] = 1,
+        bias: bool = True,
+        encoding: str = 'unroll'
+    ):
+        """Initialize PoTDepthwiseConv2d layer.
+
+        Args:
+            channels: Number of input/output channels
+            kernel_size: Size of the convolution kernel (default: 3)
+            stride: Stride of the convolution (default: 1)
+            padding: Zero-padding added to both sides (default: 1)
+            dilation: Spacing between kernel elements (default: 1)
+            bias: If True, adds a learnable bias (default: True)
+            encoding: Encoding type ('unroll', 'fp130', '5level', '2bit', 'ternary')
+        """
+        super().__init__(encoding)
+
+        self.channels = channels
+        self.in_channels = channels  # alias for compatibility
+        self.out_channels = channels  # alias for compatibility
+        self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
+        self.stride = stride if isinstance(stride, tuple) else (stride, stride)
+        self.padding = padding if isinstance(padding, tuple) else (padding, padding)
+        self.dilation = dilation if isinstance(dilation, tuple) else (dilation, dilation)
+        self.groups = channels  # depthwise: each channel is its own group
+
+        # Initialize weight parameter: [channels, 1, kH, kW]
+        self.weight = nn.Parameter(torch.empty(
+            channels, 1, *self.kernel_size
+        ))
+
+        # Initialize bias parameter
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(channels))
+        else:
+            self.register_parameter('bias', None)
+
+        # Initialize weights using Kaiming normal
+        nn.init.kaiming_normal_(self.weight, mode='fan_out', nonlinearity='relu')
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass with optional PoT quantization.
+
+        Args:
+            x: Input tensor of shape (N, C, H, W)
+
+        Returns:
+            Output tensor of shape (N, C, H_out, W_out)
+        """
+        if not self.quantize:
+            # Float mode (warmup training)
+            return F.conv2d(
+                x, self.weight, self.bias,
+                self.stride, self.padding, self.dilation, self.groups
+            )
+        
+        if self.use_integer_sim and self.scale_int is not None:
+            if self.training:
+                # Training: use float QAT for gradient flow
+                return self._forward_float_qat(x)
+            else:
+                # Eval: use integer sim for C-exact match
+                return self._forward_integer_sim(x)
+        else:
+            # Standard Float QAT Mode
+            return self._forward_float_qat(x)
+
+    def _forward_float_qat(self, x: torch.Tensor) -> torch.Tensor:
+        """Original float QAT forward."""
+        w_q = quantize_to_pot_ste(self.weight, self.alpha, encoding=self.encoding)
+
+        out = F.conv2d(
+            x, w_q * self.alpha, self.bias,
+            self.stride, self.padding, self.dilation, self.groups
+        )
+
+        if self.act_scale is not None:
+            out = quantize_activation_ste(out, self.act_scale)
+
+        return out
+
+    def _forward_integer_sim(self, x: torch.Tensor) -> torch.Tensor:
+        """Integer simulation forward - matches C inference exactly."""
+        DEBUG = False  # True로 바꾸면 상세 디버그 출력
+        
+        is_first = self.is_first_layer.item() if self.is_first_layer is not None else False
+        is_last = self.is_last_layer.item() if self.is_last_layer is not None else False
+        
+        if DEBUG:
+            print(f"\n[DEBUG DW _forward_integer_sim] is_first={is_first}, is_last={is_last}")
+            print(f"  input: shape={x.shape}, range=[{x.min():.4f}, {x.max():.4f}]")
+        
+        # Step 1: Quantize input
+        if is_first:
+            # First layer: input is NORMALIZED, denormalize with channel-wise mean
+            if self.input_mean is not None and self.input_std is not None:
+                avg_std = self.input_std.mean().item()
+                mean = self.input_mean.view(1, -1, 1, 1).to(x.device)  # [1, C, 1, 1]
+                x_raw = x * avg_std + mean  # channel-wise mean
+                x_raw = torch.clamp(x_raw, 0.0, 1.0)
+            else:
+                x_raw = x
+            # [0,1] → [0,255] (uint8), /256 absorbed in shift (+8)
+            x_int = quantize_to_uint8_ste(x_raw, 256.0)
+        else:
+            prev_scale = self.prev_act_scale if self.prev_act_scale is not None else torch.tensor(1.0)
+            x_int = quantize_to_int8_ste(x, prev_scale)
+        
+        # Step 2: PoT Depthwise Convolution with STE for gradient flow
+        w_pot = quantize_to_pot_ste(self.weight, self.alpha, encoding=self.encoding)
+        
+        acc = F.conv2d(
+            x_int, w_pot, None,
+            self.stride, self.padding, self.dilation, self.groups
+        )
+        
+        # Step 3: Requantize
+        scale_int = self.scale_int.item() if self.scale_int is not None else 1
+        shift = self.shift.item() if self.shift is not None else 0
+        acc = requantize_ste(acc, scale_int, shift)
+        
+        if DEBUG:
+            print(f"  acc after requantize: scale_int={scale_int}, shift={shift}, range=[{acc.min():.0f}, {acc.max():.0f}]")
+        
+        # Step 4: Add bias (with mean absorption for first layer)
+        if self.bias is not None:
+            act_scale = self.act_scale if self.act_scale is not None else torch.tensor(1.0)
+            
+            if is_first:
+                # First layer: absorb mean into bias
+                # Use avg_std to match QAT and C inference
+                if self.input_mean is not None and self.input_std is not None:
+                    avg_std = self.input_std.mean().item()
+                    # Depthwise: weight is [channels, 1, kH, kW]
+                    channels = w_pot.shape[0]
+                    alpha = self.alpha
+                    bias_adjusted = self.bias.clone()
+                    for c in range(channels):
+                        mean_c = self.input_mean[c].item() if c < len(self.input_mean) else 0.0
+                        weight_sum_c = w_pot[c].sum() * alpha
+                        bias_adjusted[c] = bias_adjusted[c] - (mean_c / avg_std) * weight_sum_c
+                else:
+                    bias_adjusted = self.bias
+                bias_int = round_ste(bias_adjusted * act_scale)
+                
+                if DEBUG:
+                    print(f"  [First layer DW bias absorption]")
+            else:
+                bias_int = round_ste(self.bias * act_scale)
+            
+            acc = acc + bias_int.view(1, -1, 1, 1)
+        
+        # Step 5: Clamp
+        if not is_last:
+            out = clamp_ste(acc, 0.0, 127.0)
+        else:
+            out = acc
+        
+        # Step 6: Convert back to float
+        if self.act_scale is not None and not is_last:
+            out = out / self.act_scale
+        
+        return out
+
+    def extra_repr(self) -> str:
+        """String representation of layer configuration."""
+        s = f'channels={self.channels}, kernel_size={self.kernel_size}, stride={self.stride}'
+        if self.padding != (0, 0):
+            s += f', padding={self.padding}'
+        if self.dilation != (1, 1):
+            s += f', dilation={self.dilation}'
+        if self.bias is None:
+            s += ', bias=False'
+        if self.quantize:
+            s += f', quantize=True, encoding={self.encoding}'
+        if self.use_integer_sim:
+            s += ', integer_sim=True'
+        return s

potnn/modules/linear.py

@@ -0,0 +1,199 @@
+"""PoT-quantized Linear layer with Integer Simulation.
+
+v2: Added integer simulation for C-compatible QAT
+- Forward pass can simulate C integer operations exactly
+- Matches C inference bit-for-bit when use_integer_sim=True
+- Eliminates QAT-C accuracy gap
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional
+
+from .base import PoTLayerBase
+from ..quantize.pot import quantize_to_pot_ste, quantize_to_pot, quantize_activation_ste, apply_5level_zero_constraint
+from ..quantize.integer_ops import (
+    round_half_up_ste, clamp_ste,
+    fake_quantize_input, fake_quantize_input_uint8,
+    fake_requantize
+)
+
+
+class PoTLinear(PoTLayerBase):
+    """Power-of-Two quantized Linear layer.
+
+    This layer implements a Linear (fully connected) layer with PoT weight
+    quantization.
+    
+    [Integer-Only QAT Mode]
+    The forward pass simulates C integer arithmetic EXACTLY:
+    1. Input Quantization: float -> int8 (or uint8 for first layer)
+    2. Integer Linear: int8 * int8 -> int32
+    3. Requantize: (int32 * scale_int + round) >> shift
+    4. Bias Add: + round(bias_adjusted * act_scale)
+    5. Clamp: [0, 127]
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        encoding: str = 'unroll'
+    ):
+        """Initialize PoTLinear layer.
+
+        Args:
+            in_features: Size of each input sample
+            out_features: Size of each output sample
+            bias: If True, adds a learnable bias (default: True)
+            encoding: Encoding type ('unroll', 'fp130', '5level', '2bit', 'ternary')
+        """
+        super().__init__(encoding)
+
+        self.in_features = in_features
+        self.out_features = out_features
+
+        # Initialize weight parameter
+        self.weight = nn.Parameter(torch.empty(out_features, in_features))
+
+        # Initialize bias parameter
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(out_features))
+        else:
+            self.register_parameter('bias', None)
+
+        # Initialize weights using Kaiming normal
+        nn.init.kaiming_normal_(self.weight, mode='fan_out', nonlinearity='relu')
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass with three modes:
+        1. Float warmup (quantize=False): Standard linear
+        2. Float QAT (use_integer_sim=False): PoT weight + float activation
+        3. Integer sim (use_integer_sim=True): C-identical integer ops
+        """
+        if not self.quantize:
+            # Float mode (warmup training)
+            return F.linear(x, self.weight, self.bias)
+            
+        if not getattr(self, 'use_integer_sim', False):
+            # Float QAT: PoT weight + float activation
+            # ReLU는 모델에서 외부로 호출
+            w_pot = quantize_to_pot_ste(self.weight, self.alpha, encoding=self.encoding)
+            
+            # 5level constraint
+            if self.encoding == '5level' and self.enforce_5level_constraint:
+                w_pot = apply_5level_zero_constraint(w_pot)
+            
+            out = F.linear(x, w_pot * self.alpha, self.bias)
+            return out
+        
+        # === Integer Simulation Mode (C-identical) ===
+        
+        # === 1. Prepare Integer Parameters ===
+        # Always compute dynamically to ensure consistency with export
+        scale_int, shift, _ = self._compute_scale_and_shift()
+
+        is_first = self.is_first_layer.item() if self.is_first_layer is not None else False
+        is_last = self.is_last_layer.item() if self.is_last_layer is not None else False
+        
+        # === 2. Input Quantization ===
+        if is_first:
+            # First layer: input is NORMALIZED (x - mean) / std
+            # Simulate C behavior: raw uint8 input
+            if self.input_mean is not None and self.input_std is not None:
+                avg_std = self.input_std.mean().item()
+                # Handle flattened input for Linear layer
+                if x.dim() == 2:  # [batch, features]
+                    num_ch = len(self.input_mean)
+                    feat_per_ch = x.shape[1] // num_ch
+                    if x.shape[1] == num_ch * feat_per_ch:
+                        # Reshape to apply channel-wise mean
+                        x_reshaped = x.view(x.shape[0], num_ch, feat_per_ch)
+                        mean = self.input_mean.view(1, -1, 1).to(x.device)
+                        x_raw = x_reshaped * avg_std + mean
+                        x_raw = x_raw.view(x.shape[0], -1)
+                    else:
+                        # Fallback to average mean if dimensions don't match
+                        avg_mean = self.input_mean.mean().item()
+                        x_raw = x * avg_std + avg_mean
+                else:
+                    avg_mean = self.input_mean.mean().item()
+                    x_raw = x * avg_std + avg_mean
+                
+                x_raw = clamp_ste(x_raw, 0.0, 1.0)
+            else:
+                x_raw = x
+                
+            # Quantize to uint8 [0, 255]
+            x_int = fake_quantize_input_uint8(x_raw, 256.0)
+        else:
+            # Other layers: Input is already int8 from previous layer
+            x_int = x
+
+        # === 3. Weight Quantization ===
+        w_pot = quantize_to_pot_ste(self.weight, self.alpha, encoding=self.encoding)
+        
+        # 5level constraint (always apply for 5level encoding to match export)
+        if self.encoding == '5level':
+            w_pot = apply_5level_zero_constraint(w_pot)
+
+        # === 4. Integer Linear ===
+        # F.linear with integer-valued inputs/weights -> integer-valued output
+        acc = F.linear(x_int, w_pot, None)
+
+        # === 5. Requantize ===
+        acc_scaled = fake_requantize(acc, scale_int, shift)
+
+        # === 6. Bias Addition ===
+        if self.bias is not None:
+            act_scale = self.act_scale if self.act_scale is not None else torch.tensor(1.0)
+            
+            if is_first and self.input_mean is not None and self.input_std is not None:
+                # Absorb mean/std into bias
+                avg_std = self.input_std.mean().item()
+                alpha = self.alpha
+                in_features = self.weight.shape[1]
+                bias_adjusted = self.bias.clone()
+                
+                num_channels = len(self.input_mean)
+                features_per_channel = in_features // num_channels if num_channels > 0 else in_features
+                
+                if num_channels > 0 and in_features == num_channels * features_per_channel:
+                    for c in range(num_channels):
+                        mean_c = self.input_mean[c].item()
+                        start_idx = c * features_per_channel
+                        end_idx = start_idx + features_per_channel
+                        weight_sum_c = w_pot[:, start_idx:end_idx].sum(dim=1) * alpha
+                        bias_adjusted = bias_adjusted - (mean_c / avg_std) * weight_sum_c
+                else:
+                    # Fallback if dimensions don't match
+                    bias_adjusted = self.bias
+            else:
+                bias_adjusted = self.bias
+
+            # Quantize bias: round(bias * act_scale)
+            bias_int = round_half_up_ste(bias_adjusted * act_scale)
+            
+            # Add bias
+            acc_scaled = acc_scaled + bias_int
+
+        # === 7. Clamp (ReLU) ===
+        if not is_last:
+            out = clamp_ste(acc_scaled, 0.0, 127.0)
+        else:
+            out = acc_scaled
+
+        # === 8. Output ===
+        # Round to ensure exact integer (floating point precision)
+        # Use STE to maintain gradient flow during training
+        # int8 그대로 반환 (C와 동일)
+        out = round_half_up_ste(out)
+
+        return out
+
+    def extra_repr(self) -> str:
+        s = super().extra_repr()
+        s += f', quantize={self.quantize}'
+        return s

potnn/quantize/__init__.py

@@ -0,0 +1,35 @@
+"""Quantization module for potnn."""
+
+from .pot import (
+    quantize_to_pot, quantize_to_pot_ste, quantize_activation_ste, 
+    quantize_activation_c_aligned_ste, get_pot_values
+)
+from .calibration import calibrate_model
+from .qat import prepare_qat, alpha_reg_loss, enable_integer_sim, disable_integer_sim
+from .integer_sim import (
+    round_ste, round_half_up_ste, floor_ste, clamp_ste,
+    quantize_to_int8_ste, quantize_to_uint8_ste,
+    requantize_ste, compute_scale_params
+)
+
+__all__ = [
+    'quantize_to_pot',
+    'quantize_to_pot_ste',
+    'quantize_activation_ste',
+    'quantize_activation_c_aligned_ste',
+    'get_pot_values',
+    'calibrate_model',
+    'prepare_qat',
+    'alpha_reg_loss',
+    'enable_integer_sim',
+    'disable_integer_sim',
+    # Integer simulation
+    'round_ste',
+    'round_half_up_ste',
+    'floor_ste', 
+    'clamp_ste',
+    'quantize_to_int8_ste',
+    'quantize_to_uint8_ste',
+    'requantize_ste',
+    'compute_scale_params',
+]

potnn/quantize/calibration.py

@@ -0,0 +1,233 @@
+"""Activation calibration for PoT quantization."""
+
+import torch
+import math
+from torch.utils.data import DataLoader
+from typing import Dict, Any, Optional, Union, List
+
+from ..modules.base import PoTLayerBase
+
+
+def calibrate_model(
+    model: torch.nn.Module,
+    data_loader: DataLoader,
+    num_batches: int = 10,
+    mean: Union[float, List[float]] = 0.0,
+    std: Union[float, List[float]] = 1.0
+) -> Dict[str, float]:
+    """Calibrate activation scales for each PoT layer.
+
+    This function measures the maximum activation values for each layer
+    during inference on calibration data. These values are used to set
+    fixed activation scales for quantization-aware training.
+
+    IMPORTANT: This function should be called ONCE before QAT training,
+    after float warmup training.
+
+    Args:
+        model: Model containing PoT layers to calibrate
+        data_loader: DataLoader for calibration data
+        num_batches: Number of batches to use for calibration (default: 10)
+        mean: Dataset mean for normalization (default: 0.0)
+        std: Dataset std for normalization (default: 1.0)
+
+    Returns:
+        Dictionary mapping layer names to maximum activation values
+    """
+    # ========================================
+    # Step 0: Fuse BatchNorm layers BEFORE calibration
+    # This ensures calibration measures activations with final fused weights.
+    # ========================================
+    from ..fuse import fuse_batchnorm
+    fuse_batchnorm(model)
+    
+    print("Starting calibration...")
+
+    # Set model to evaluation mode
+    model.eval()
+
+    # Dictionary to store max activation values for each layer
+    activation_max = {}
+    hooks = []
+
+    def make_hook(name):
+        """Create a forward hook to capture activation values."""
+        def hook(module, input, output):
+            with torch.no_grad():
+                max_val = output.abs().max().item()
+                if name not in activation_max:
+                    activation_max[name] = max_val
+                else:
+                    activation_max[name] = max(activation_max[name], max_val)
+        return hook
+
+    # Register hooks on all PoT layers
+    for name, module in model.named_modules():
+        if isinstance(module, PoTLayerBase):
+            hook = module.register_forward_hook(make_hook(name))
+            hooks.append(hook)
+    
+    # Also register hooks on PoTAdd layers to track their output range
+    from ..modules.add import PoTAdd
+    
+    # PoTAdd 입력 scale 추적용 딕셔너리
+    add_input_max = {}  # {name: {'x': max_x, 'y': max_y}}
+    
+    def make_add_hook(name):
+        """Create a forward hook for PoTAdd to capture input/output values."""
+        def hook(module, input, output):
+            with torch.no_grad():
+                # output max
+                max_val = output.abs().max().item()
+                if name not in activation_max:
+                    activation_max[name] = max_val
+                else:
+                    activation_max[name] = max(activation_max[name], max_val)
+                
+                # input max (x=skip, y=conv)
+                x, y = input[0], input[1]
+                max_x = x.abs().max().item()
+                max_y = y.abs().max().item()
+                
+                if name not in add_input_max:
+                    add_input_max[name] = {'x': max_x, 'y': max_y}
+                else:
+                    add_input_max[name]['x'] = max(add_input_max[name]['x'], max_x)
+                    add_input_max[name]['y'] = max(add_input_max[name]['y'], max_y)
+        return hook
+    
+    for name, module in model.named_modules():
+        if isinstance(module, PoTAdd):
+            hook = module.register_forward_hook(make_add_hook(name))
+            hooks.append(hook)
+
+    # Run forward passes to collect statistics
+    with torch.no_grad():
+        for i, batch in enumerate(data_loader):
+            if i >= num_batches:
+                break
+
+            # Handle different data formats (data, target) or just data
+            if isinstance(batch, (list, tuple)):
+                data = batch[0]
+            else:
+                data = batch
+
+            # Move to same device as model
+            device = next(model.parameters()).device
+            data = data.to(device)
+            
+            # Normalize data (support both scalar and per-channel mean/std)
+            # Dynamic reshape based on data dimensions (3D for Conv1d, 4D for Conv2d)
+            if isinstance(mean, (list, tuple)):
+                mean_t = torch.tensor(mean, dtype=data.dtype, device=device)
+                if data.dim() == 3:  # Conv1d: (B, C, L)
+                    mean_t = mean_t.view(1, -1, 1)
+                else:  # Conv2d: (B, C, H, W)
+                    mean_t = mean_t.view(1, -1, 1, 1)
+            else:
+                mean_t = mean
+            if isinstance(std, (list, tuple)):
+                std_t = torch.tensor(std, dtype=data.dtype, device=device)
+                if data.dim() == 3:  # Conv1d: (B, C, L)
+                    std_t = std_t.view(1, -1, 1)
+                else:  # Conv2d: (B, C, H, W)
+                    std_t = std_t.view(1, -1, 1, 1)
+            else:
+                std_t = std
+            data = (data - mean_t) / std_t
+
+            # Forward pass
+            _ = model(data)
+
+            # Progress indicator
+            if (i + 1) % 5 == 0:
+                print(f"  Calibration batch {i + 1}/{min(num_batches, len(data_loader))}")
+
+    # Remove hooks
+    for hook in hooks:
+        hook.remove()
+
+    # Set activation scales based on calibration
+    print("\nSetting activation scales:")
+    for name, module in model.named_modules():
+        if isinstance(module, PoTLayerBase):
+            if name in activation_max and activation_max[name] > 0:
+                # Set the activation scale for quantization
+                module.calibrate(activation_max[name])
+                print(f"  {name}: act_max={activation_max[name]:.2f}, "
+                      f"act_scale={module.act_scale.item():.4f}")
+            else:
+                # Default if layer wasn't activated during calibration
+                module.calibrate(1.0)
+                print(f"  {name}: no activations detected, using default scale")
+    
+    # Set activation scales for PoTAdd layers
+    print("\nSetting PoTAdd scales:")
+    for name, module in model.named_modules():
+        if isinstance(module, PoTAdd):
+            if name in activation_max and activation_max[name] > 0:
+                # act_scale = 127 / max_activation
+                act_scale = 127.0 / activation_max[name]
+                module.act_scale = torch.tensor(act_scale)
+                module.quantize = True
+                
+                # 입력 scale 설정 (핵심!)
+                if name in add_input_max:
+                    max_x = add_input_max[name]['x']
+                    max_y = add_input_max[name]['y']
+                    scale_x = 127.0 / max_x if max_x > 0 else 1.0
+                    scale_y = 127.0 / max_y if max_y > 0 else 1.0
+                    module.set_scales(scale_x, scale_y)
+                    print(f"  {name}: scale_x={scale_x:.4f}, scale_y={scale_y:.4f}, "
+                          f"ratio={scale_y/scale_x:.4f} (PoTAdd)")
+                else:
+                    print(f"  {name}: act_max={activation_max[name]:.2f}, "
+                          f"act_scale={act_scale:.4f} (PoTAdd, no input scales)")
+            else:
+                # Default scale
+                module.act_scale = torch.tensor(1.0)
+                module.quantize = True
+                print(f"  {name}: no activations detected, using default scale (PoTAdd)")
+
+    # Initialize alpha values based on weight distribution
+    print("\nInitializing alpha values based on weight distribution:")
+    for name, module in model.named_modules():
+        if isinstance(module, PoTLayerBase):
+            with torch.no_grad():
+                # Calculate weight standard deviation
+                w_std = module.weight.std().item()
+
+                # Calculate appropriate raw_alpha value
+                # We want: softplus(raw_alpha) ≈ w_std
+                # softplus(x) = log(1 + exp(x))
+                # So we need: x = log(exp(w_std) - 1) when w_std > log(2)
+                if w_std > 0.01:
+                    # Inverse softplus: raw = log(exp(target) - 1)
+                    # But ensure numerical stability
+                    target = w_std
+                    if target > 10:  # Avoid overflow
+                        raw = target  # For large values, softplus(x) ≈ x
+                    elif target > 0.1:
+                        raw = math.log(math.exp(target) - 1)
+                    else:
+                        # For small values, use approximation
+                        raw = math.log(target) if target > 0 else -2.0
+                else:
+                    # Very small weights, use small alpha
+                    raw = -2.0
+
+                # Update raw_alpha
+                module.raw_alpha.data.fill_(raw)
+
+                # Update alpha_init to match the newly initialized alpha
+                # This ensures regularization pulls towards the calibrated value
+                module.alpha_init.data.fill_(module.alpha.item())
+
+                # Verify the result
+                actual_alpha = module.alpha.item()
+                print(f"  {name}: weight_std={w_std:.4f}, "
+                      f"raw_alpha={raw:.4f}, alpha={actual_alpha:.4f}")
+
+    print(f"\nCalibration complete. Processed {len(activation_max)} layers.")
+    return activation_max