SparseMIMOStem.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import spconv.pytorch as spconv                 # ← new
from detectron2.modeling.backbone.resnet import CNNBlockBase
import fvcore.nn.weight_init as weight_init

class SparseMIMOStem(CNNBlockBase):
    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size=(1, 12),
        dilation=(1, 16),
        use_bn=False,
        padding_ants=96,
        stride=1,
        norm="",
    ):
        super().__init__(in_channels, out_channels, stride)
        self.padding_ants = padding_ants
        self.use_bn = use_bn

        # one sub-manifold 2-D sparse convolution
        # compute symmetric padding for height _and_ width
        pad_h = (kernel_size[0] - 1) * dilation[0] // 2
        pad_w = (kernel_size[1] - 1) * dilation[1] // 2
        self.sconv = spconv.SubMConv2d(
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            dilation=dilation,
            stride=(1, 1),
            padding=(pad_h, pad_w),       # now pads both dims correctly
            bias=not use_bn,
        )
        weight_init.c2_msra_fill(self.sconv)
        if use_bn:
            self.bn = nn.BatchNorm1d(out_channels)  # acts on sparse features

        # -------- FLOP bookkeeping -------------------------------------
        # one buffer to hold *this* module’s FLOPs from the most recent fwd
        self.register_buffer("_flops", torch.zeros(1, dtype=torch.long),
                             persistent=False)       # persistent=False → not saved in checkpoints


    """     # ---------- helper to build a SparseConvTensor ---------- default
    @staticmethod
    def _dense_to_sparse(x):
        
        #Args:
            #x: (B,C,H,W) dense float tensor
        #Returns:
            #st: spconv.SparseConvTensor with the same feature dtype
        
        B, C, H, W = x.shape
        device = x.device

        # mask all-zero pixels (across channels)
        nz_mask = x.view(B, C, -1).abs().sum(1) != 0          # (B, H*W)
        idxs, feats = [], []

        for b in range(B):

            # get the flat (H*W,) mask for batch b
            m = nz_mask[b]                        # shape (H*W,)
            # reshape back to (H, W)
            m2 = m.view(H, W)                     # shape (H, W)
            # now nonzero returns (row, col) pairs directly
            ys, xs = m2.nonzero(as_tuple=True)   # each (N_points,)
            if ys.numel() == 0:                   # skip if no active pixels
                continue
            # build (batch_index, y, x) coords
            coords = torch.stack([
                torch.full_like(ys, b),  # batch dim
                ys,                     # row indices
                xs                      # col indices
            ], dim=1)                     # shape (N_points, 3)
            idxs.append(coords)
            # gather features at those y,x locations
            feats.append(x[b, :, ys, xs].t())  # (N_points, C)


        idxs = torch.cat(idxs, 0).int()       # (N, 3)  (b, y, x)
        feats = torch.cat(feats, 0).contiguous()  # (N, C)

        return spconv.SparseConvTensor(
            features=feats,
            indices=idxs,
            spatial_shape=(H, W),
            batch_size=B,
        )

    # --------------------------------------------------------- """

    """     @staticmethod
    def _dense_to_sparse(x, threshold: float = 0.000):
        
        #Convert a dense tensor to a SparseConvTensor, dropping anything
        #whose absolute‐sum across channels is <= threshold.

        #Args:
            #x: Tensor of shape (B, C, H, W)
            #threshold: any location whose sum(abs(x), dim=1) <= threshold
                       #will be treated as zero (and dropped)

        #Returns:
            #st: spconv.SparseConvTensor with
                #features shape (N, C) and indices shape (N, 3),
                #where N = #active positions.
        
        B, C, H, W = x.shape
        # 1) zero out “small” values in‐place (optional)
        if threshold > 0:
            x = x.clone()
            mask_vals = x.abs() <= threshold
            print(f"[dense_to_sparse] zeroed {mask_vals.sum().item()} / {x.numel()} entries (thr={threshold})")
            x[mask_vals] = 0.0

        # 2) build a (B, H, W) boolean mask of “active” sites
        #    (sum across channels > 0)
        m = x.abs().sum(dim=1) > 0   # shape: (B, H, W)

        # DEBUG: how many “active” pixels remain?
        total_positions = B * H * W
        active = int(m.sum().item())
        print(f"[dense_to_sparse] keeping {active} / {total_positions} sites")

        # 3) get the (batch, y, x) coordinates of all True entries
        b_inds, ys, xs = m.nonzero(as_tuple=True)  # each is (N,)

        # if nothing is active, return an empty tensor
        if b_inds.numel() == 0:
            return spconv.SparseConvTensor(
                features=torch.zeros((0, C), device=x.device),
                indices=torch.zeros((0, 3), dtype=torch.int32, device=x.device),
                spatial_shape=(H, W),
                batch_size=B,
            )

        # 4) gather the per-point features: (N, C)
        feats = x[b_inds, :, ys, xs]  # fancy‐indexing yields (N, C)

        # 5) stack coords into (N, 3) int tensor
        coords = torch.stack([b_inds, ys, xs], dim=1).int()

        return spconv.SparseConvTensor(
            features=feats,
            indices=coords,
            spatial_shape=(H, W),
            batch_size=B,
        )
    """

    """     @staticmethod
    def _dense_to_sparse(x, threshold: float = 0.0):
        
        #Convert a dense tensor to SparseConvTensor, dropping *only* those
        #(b,y,x) positions where sum(abs(x), dim=1) == 0 (or <= threshold).
        
        B, C, H, W = x.shape

        # 1) Compute the per‐location absolute sum over channels
        abs_sum = x.abs().sum(dim=1)   # shape (B, H, W)

        # DEBUG: see the min/max of abs_sum to confirm what values you have
        print(
            f"[dense_to_sparse] abs_sum.min={abs_sum.min().item():.3e}, "
            f"max={abs_sum.max().item():.3e}"
        )

        # 2) Build mask directly on that sum
        #    Keep positions where abs_sum > threshold
        m = abs_sum > threshold

        # DEBUG: how many sites survive?
        total_sites = B * H * W
        kept = int(m.sum().item())
        print(f"[dense_to_sparse] keeping {kept} / {total_sites} sites, thr={threshold}")

        # 3) Get coordinates of the kept sites
        b_inds, ys, xs = m.nonzero(as_tuple=True)  # each has length N

        # 4) If none, return an empty SparseConvTensor
        if b_inds.numel() == 0:
            return spconv.SparseConvTensor(
                features=torch.zeros((0, C), device=x.device),
                indices=torch.zeros((0, 3), dtype=torch.int32, device=x.device),
                spatial_shape=(H, W),
                batch_size=B,
            )

        # 5) Gather features & build coords
        feats  = x[b_inds, :, ys, xs]                    # (N, C)
        coords = torch.stack([b_inds, ys, xs], dim=1).int()

        return spconv.SparseConvTensor(
            features=feats,
            indices=coords,
            spatial_shape=(H, W),
            batch_size=B,
        ) """

    """     def forward(self, x):
        
        #x: dense (B,C,H,W)
        
        B, C, H, W = x.shape

        # 1) circular padding (wrap) on width
        x = torch.cat(
            [x[..., -self.padding_ants:], x, x[..., :self.padding_ants]], dim=3
        )  # new W = W + 2*padding_ants

        # 2) dense → sparse
        #st = self._dense_to_sparse(x)
        st = self._dense_to_sparse(x, threshold=0.0)

        # 3) sparse convolution (features stay sparse)
        st = self.sconv(st)

        # 4) optional BN + ReLU: replace the feature tensor, not by setting .features
        feats = st.features
        if self.use_bn:
            feats = self.bn(feats)
        feats = F.relu(feats)
        st = st.replace_feature(feats)

        # 5) back to dense so Detectron2 can keep going
        x_dense = st.dense()                  # (B, C_out, H', W')
        # 6) crop width back to original W
        left = (x_dense.shape[-1] - W) // 2
        x_dense = x_dense[..., left : left + W]

        return x_dense """
    @staticmethod
    def _dense_to_sparse(x_wrapped, threshold: float = 0.0, mask_source=None):
        """
        Args:
            x_wrapped: wrapped tensor (B, C, H, W_new) to extract features from
            threshold: threshold for sparsity detection (should be 0.0)
            mask_source: optional tensor (B, C, H, W_orig) used to compute sparsity mask
        """
        B, C, H, W_wrapped = x_wrapped.shape

        if mask_source is None:
            raise ValueError("mask_source must be provided to detect sparsity before wrapping")

        _, _, _, W_orig = mask_source.shape
        abs_sum = mask_source.abs().sum(dim=1)  # (B, H, W_orig)
        
        # Print stats about abs_sum before thresholding
        #print(f"[dense_to_sparse] abs_sum.min={abs_sum.min():.3e}, max={abs_sum.max():.3e}")

        # Create mask of active locations
        mask = abs_sum > threshold  # (B, H, W_orig)

        # Count how many entries were zeroed (not active)
        #total_sites = B * H * W_orig
        # = int(mask.sum().item())
        #num_zeroed = total_sites - num_active
        #print(f"[dense_to_sparse] active sites: {num_active} / {total_sites} (zeroed: {num_zeroed})")

        # Embed mask into the wrapped tensor
        pad = (W_wrapped - W_orig) // 2
        mask_full = torch.zeros((B, H, W_wrapped), dtype=torch.bool, device=x_wrapped.device)
        mask_full[..., pad:pad + W_orig] = mask

        # Get sparse coordinates
        b_inds, ys, xs = mask_full.nonzero(as_tuple=True)
        if b_inds.numel() == 0:
            #print("[dense_to_sparse] WARNING: no active elements found.")
            return spconv.SparseConvTensor(
                features=torch.zeros((0, C), device=x_wrapped.device),
                indices=torch.zeros((0, 3), dtype=torch.int32, device=x_wrapped.device),
                spatial_shape=(H, W_wrapped),
                batch_size=B,
            )

        # Gather features and indices
        feats = x_wrapped[b_inds, :, ys, xs]
        coords = torch.stack([b_inds, ys, xs], dim=1).int()

        #print(f"[dense_to_sparse] SparseConvTensor: features={feats.shape}, indices={coords.shape}")
        return spconv.SparseConvTensor(
            features=feats,
            indices=coords,
            spatial_shape=(H, W_wrapped),
            batch_size=B,
        )

    def forward(self, x):
        """
        x: dense (B,C,H,W)
        """
        B, C, H, W = x.shape

        # === [1] Threshold original (unwrapped) data BEFORE padding ===
        # We'll use this for sparse site detection.
        # Clone to avoid modifying input
        x_thresh = x.clone()
        x_thresh[x_thresh.abs() < 1e-4] = 0.0
        # DEBUG: count zeros in x_thresh
        #abs_sum = x_thresh.abs().sum(dim=1)  # shape (B, H, W)
        #num_zero_locations = (abs_sum == 0).sum().item()
        #total_locations = B * H * W
        #print(f"[forward] x_thresh zero locations: {num_zero_locations} / {total_locations}")

        # === [2] Now do the circular wrap on the ORIGINAL x (not x_thresh) ===
        x_wrapped = torch.cat(
            [x[..., -self.padding_ants:], x, x[..., :self.padding_ants]], dim=3
        )  # shape: (B, C, H, W + 2*padding)

        # === [3] Convert to sparse, using original sparsity pattern ===
        # We use x_thresh to determine which coordinates are active,
        # but we use x_wrapped for the actual feature values.
        st = self._dense_to_sparse(x_wrapped, threshold=0.0, mask_source=x_thresh)

        # === [4] Sparse convolution ===
        st = self.sconv(st)

        # === [5] BN + ReLU ===
        feats = st.features
        if self.use_bn:
            feats = self.bn(feats)
        feats = F.relu(feats)
        st = st.replace_feature(feats)

        # === [6] Convert back to dense and crop to original W ===
        x_dense = st.dense()  # shape: (B, C_out, H, W + pad_crop)
        left = (x_dense.shape[-1] - W) // 2
        x_dense = x_dense[..., left : left + W]

        # -------- FLOP *after* convolution ------------------------------
        #  #active sites == st.indices.shape[0]
        n_active = st.indices.shape[0]
        k = self.sconv.kernel_size[0] * self.sconv.kernel_size[1]
        cin, cout = self.sconv.in_channels, self.sconv.out_channels
        macs = n_active * k * cin * cout            # multiply–accumulates
        flops_here = macs * 2                      # 1 mul + 1 add per MAC

        return x_dense