Upload SoftPool.py with huggingface_hub

SoftPool.py

@@ -0,0 +1,80 @@
+"""
+Pure PyTorch implementation of SoftPool.
+This is a fallback that doesn't require CUDA kernel compilation.
+SoftPool: https://arxiv.org/abs/2101.00440
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def soft_pool2d(x, kernel_size=(2, 2), stride=None, force_inplace=False):
+    """
+    Apply soft pooling on 2D input tensor.
+    
+    SoftPool approximates max pooling while maintaining differentiability
+    by using exponential weighting: y = sum(x * exp(x)) / sum(exp(x))
+    
+    Args:
+        x: Input tensor of shape (N, C, H, W)
+        kernel_size: Pooling kernel size
+        stride: Stride (defaults to kernel_size)
+        force_inplace: Unused, for API compatibility
+    
+    Returns:
+        Pooled tensor
+    """
+    if stride is None:
+        stride = kernel_size
+    
+    if isinstance(kernel_size, int):
+        kernel_size = (kernel_size, kernel_size)
+    if isinstance(stride, int):
+        stride = (stride, stride)
+    
+    # Use unfold to extract patches
+    batch, channels, height, width = x.shape
+    kh, kw = kernel_size
+    sh, sw = stride
+    
+    # Calculate output dimensions
+    out_h = (height - kh) // sh + 1
+    out_w = (width - kw) // sw + 1
+    
+    # Apply exponential weighting
+    # For numerical stability, subtract max before exp
+    x_unfold = F.unfold(x, kernel_size=kernel_size, stride=stride)  # (N, C*kh*kw, out_h*out_w)
+    x_unfold = x_unfold.view(batch, channels, kh * kw, out_h * out_w)
+    
+    # Softmax-style weighting for soft pooling
+    x_max = x_unfold.max(dim=2, keepdim=True)[0]
+    exp_x = torch.exp(x_unfold - x_max)  # Numerical stability
+    
+    # Weighted sum: sum(x * exp(x)) / sum(exp(x))
+    softpool = (x_unfold * exp_x).sum(dim=2) / (exp_x.sum(dim=2) + 1e-8)
+    
+    # Reshape to output format
+    softpool = softpool.view(batch, channels, out_h, out_w)
+    
+    return softpool
+
+
+class SoftPool2d(nn.Module):
+    """
+    SoftPool 2D Layer.
+    
+    A differentiable pooling operation that approximates max pooling
+    using exponential weighting.
+    """
+    
+    def __init__(self, kernel_size=(2, 2), stride=None, force_inplace=False):
+        super(SoftPool2d, self).__init__()
+        self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
+        self.stride = stride if stride is not None else self.kernel_size
+        self.force_inplace = force_inplace
+    
+    def forward(self, x):
+        return soft_pool2d(x, self.kernel_size, self.stride, self.force_inplace)
+    
+    def extra_repr(self):
+        return f'kernel_size={self.kernel_size}, stride={self.stride}'