1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_pool.py",
4  "start_line": 31,
5  "end_line": 48,
6  "content": "class CustomPool(nn.Module):\n    \"\"\"Spatial-Channel Attention Pooling (SCAP).\n\n    Learns adaptive spatial attention weights per channel group,\n    combining attentive averaging with max pooling via learnable gating.\n    \"\"\"\n\n    def __init__(self):\n        super().__init__()\n        # Will be initialized in forward when input shape is known\n        self.spatial_att = None\n        self.gate_weight = nn.Parameter(torch.tensor(0.5))\n        self.gate_bias = nn.Parameter(torch.tensor(0.0))\n\n    def _create_spatial_attention(self, x):\n        B, C, H, W = x.shape\n        # Use a lightweight network to compute spatial attention\n        # First reduce channels via group convolution, then produce HxW attention map\n        # We'll create a small network: Conv2d with kernel 1 to reduce to C//8, then another to 1 channel\n        # Use adaptive architecture based on input size\n        mid_channels = max(C // 8, 8)\n        self.spatial_att = nn.Sequential(\n            nn.Conv2d(C, mid_channels, kernel_size=1, bias=False),\n            nn.BatchNorm2d(mid_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(mid_channels, 1, kernel_size=1, bias=False),\n            nn.Sigmoid()\n        ).to(x.device)\n\n    def forward(self, x):\n        B, C, H, W = x.shape\n        if self.spatial_att is None:\n            self._create_spatial_attention(x)\n        \n        # Compute spatial attention weights\n        att_weights = self.spatial_att(x)  # [B, 1, H, W]\n        \n        # Weighted average pooling using attention\n        weighted_sum = (x * att_weights).sum(dim=(2, 3))  # [B, C]\n        weights_sum = att_weights.sum(dim=(2, 3))  # [B, 1]\n        attended_avg = weighted_sum / (weights_sum + 1e-6)  # [B, C]\n        \n        # Global max pooling\n        global_max = F.adaptive_max_pool2d(x, 1).view(B, C)\n        \n        # Learnable gating between attended average and max\n        gate = torch.sigmoid(self.gate_weight * 2.0 + self.gate_bias)\n        out = gate * attended_avg + (1 - gate) * global_max\n        return out"
7}

1{
2  "n": 1
3}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_pool.py",
4  "start_line": 31,
5  "end_line": 48,
6  "content": "class CustomPool(nn.Module):\n    \"\"\"Channel-wise Adaptive Pooling (CAP).\n\n    Learns per-channel gating between average and max pooling,\n    allowing each channel to adapt its aggregation strategy.\n    Also incorporates channel statistics (mean, max, std) to compute gates.\n    \"\"\"\n\n    def __init__(self):\n        super().__init__()\n        # Statistics network: from 3 stats per channel to gate weight\n        # Input: 3 statistics (mean, max, std) per channel\n        # Output: gate weight per channel\n        self.stat_net = nn.Sequential(\n            nn.Linear(3, 16),\n            nn.ReLU(inplace=True),\n            nn.Linear(16, 1, bias=False),\n        )\n        # Learnable bias per channel\n        self.gate_bias = nn.Parameter(torch.zeros(1))\n        # Learnable per-channel scaling factor\n        self.channel_scale = nn.Parameter(torch.ones(1))\n\n    def forward(self, x):\n        B, C, H, W = x.shape\n        # Compute basic statistics per channel\n        x_flat = x.view(B, C, -1)  # [B, C, H*W]\n        mean = x_flat.mean(dim=2)  # [B, C]\n        maximum = x_flat.max(dim=2).values  # [B, C]\n        std = x_flat.std(dim=2)  # [B, C]\n        \n        # Concatenate statistics along new dimension for stat_net\n        stats = torch.stack([mean, maximum, std], dim=2)  # [B, C, 3]\n        # Process each channel independently\n        # Reshape to [B*C, 3], apply stat_net, reshape back\n        stats_flat = stats.view(-1, 3)\n        gate_raw = self.stat_net(stats_flat).view(B, C)  # [B, C]\n        gate = torch.sigmoid(gate_raw * self.channel_scale + self.gate_bias)\n        \n        # Compute average and max pooling\n        avg_pool = F.adaptive_avg_pool2d(x, 1).view(B, C)\n        max_pool = F.adaptive_max_pool2d(x, 1).view(B, C)\n        \n        # Channel-wise gated combination\n        out = gate * avg_pool + (1 - gate) * max_pool\n        return out"
7}

1{
2  "n": 1
3}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_pool.py",
4  "start_line": 31,
5  "end_line": 76,
6  "content": "class CustomPool(nn.Module):\n    \"\"\"Learnable Channel-wise Gated Pooling (LCGP).\n\n    Learns per-channel gating weight between average and max pooling.\n    Each channel can adapt its own aggregation strategy.\n    Simple, stable, and effective across architectures.\n    \"\"\"\n\n    def __init__(self):\n        super().__init__()\n        # Per-channel gate weight (unscaled). Shape will be set in forward.\n        self.gate_weight = None\n        # Temperature parameter (learnable)\n        self.temperature = nn.Parameter(torch.tensor(1.0))\n        # Small epsilon for numerical stability\n        self.eps = 1e-6\n\n    def _initialize_gate_weight(self, C):\n        # Initialize gate weight for C channels\n        # Initialize to zero so sigmoid(gate_weight) = 0.5 (equal blend)\n        self.gate_weight = nn.Parameter(torch.zeros(C))\n\n    def forward(self, x):\n        B, C, H, W = x.shape\n        if self.gate_weight is None or len(self.gate_weight) != C:\n            self._initialize_gate_weight(C)\n        \n        # Compute average and max pooling\n        avg_pool = F.adaptive_avg_pool2d(x, 1).view(B, C)\n        max_pool = F.adaptive_max_pool2d(x, 1).view(B, C)\n        \n        # Per-channel gate\n        gate = torch.sigmoid(self.gate_weight * self.temperature)\n        # Ensure gate is broadcastable across batch dimension\n        gate = gate.view(1, C)\n        \n        # Gated combination\n        out = gate * avg_pool + (1 - gate) * max_pool\n        return out"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_pool.py",
4  "start_line": 31,
5  "end_line": 69,
6  "content": "class CustomPool(nn.Module):\n    \"\"\"Learnable Channel-wise Gated Pooling (LCGP).\n\n    Learns per-channel gating weight between average and max pooling.\n    Each channel can adapt its own aggregation strategy.\n    Simple, stable, and effective across architectures.\n    \"\"\"\n\n    def __init__(self):\n        super().__init__()\n        # Maximum number of channels expected across architectures\n        # MobileNetV2 has 1280 channels, we add margin\n        self.max_channels = 2048\n        # Per-channel gate weight (unscaled). Initialize to zero -> sigmoid(0)=0.5\n        self.gate_weight = nn.Parameter(torch.zeros(self.max_channels))\n        # Temperature parameter (learnable)\n        self.temperature = nn.Parameter(torch.tensor(1.0))\n        # Small epsilon for numerical stability\n        self.eps = 1e-6\n\n    def forward(self, x):\n        B, C, H, W = x.shape\n        if C > self.max_channels:\n            # Expand gate_weight if needed (unlikely)\n            new_weight = torch.zeros(C, device=x.device)\n            new_weight[:self.max_channels] = self.gate_weight\n            self.gate_weight = nn.Parameter(new_weight)\n            self.max_channels = C\n        \n        # Slice gate_weight to current channel count\n        gate_w = self.gate_weight[:C]\n        \n        # Compute average and max pooling\n        avg_pool = F.adaptive_avg_pool2d(x, 1).view(B, C)\n        max_pool = F.adaptive_max_pool2d(x, 1).view(B, C)\n        \n        # Per-channel gate\n        gate = torch.sigmoid(gate_w * self.temperature)\n        # Ensure gate is broadcastable across batch dimension\n        gate = gate.view(1, C)\n        \n        # Gated combination\n        out = gate * avg_pool + (1 - gate) * max_pool\n        return out"
7}