custom_backboneSparseA1.py

import torch.nn as nn
import torch
from detectron2.modeling.backbone import build_resnet_backbone, Backbone, build_backbone
from detectron2.modeling import BACKBONE_REGISTRY
from detectron2.modeling.backbone.fpn import build_resnet_fpn_backbone
from detectron2.modeling.backbone.fpn import FPN
from detectron2.layers import ShapeSpec
from detectron2.modeling.backbone import Backbone
import spconv.pytorch as spconv
from detectron2.modeling.backbone.resnet import BottleneckBlock as _BottleneckBlock, CNNBlockBase
from detectron2.modeling.backbone.resnet import get_norm, Conv2d
import fvcore.nn.weight_init as weight_init
import detectron2.modeling.backbone.resnet as _resnet_module
import torch.nn.functional as F

class SparseBottleneckBlock(CNNBlockBase):
    
    #A drop-in replacement for Detectron2's BottleneckBlock,
    #but uses SubMConv2d in place of Conv2d.
    

    def __init__(
        self,
        in_channels,
        out_channels,
        *,
        bottleneck_channels,
        stride=1,
        num_groups=1,
        norm="BN",
        stride_in_1x1=False,
        dilation=1,
    ):
        super().__init__(in_channels, out_channels, stride)

        """         # 1) Shortcut conv if needed
        if in_channels != out_channels:
            self.shortcut_conv = spconv.SubMConv2d(
                in_channels, out_channels,
                kernel_size=1, stride=stride,
                padding=0, dilation=1, bias=False,
            )
            self.shortcut_bn = nn.BatchNorm1d(out_channels)
        else:
            self.shortcut_conv = None """
        
        # 1) Shortcut conv if needed
        if in_channels != out_channels:
            ConvClass = spconv.SparseConv2d if stride > 1 else spconv.SubMConv2d
            self.shortcut_conv = ConvClass(
                in_channels, out_channels,
                kernel_size=1, stride=stride,
                padding=0, dilation=1, bias=False,
            )
            self.shortcut_bn = nn.BatchNorm1d(out_channels)
        else:
            self.shortcut_conv = None

        # 2) Decide stride placement
        stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)

        # 3) conv1: 1×1
        """         self.conv1 = spconv.SubMConv2d(
            in_channels, bottleneck_channels,
            kernel_size=1, stride=stride_1x1,
            padding=0, dilation=1, bias=False,
        ) """
        Conv1 = spconv.SparseConv2d if stride_1x1 > 1 else spconv.SubMConv2d
        self.conv1 = Conv1(
            in_channels, bottleneck_channels,
            kernel_size=1, stride=stride_1x1,
            padding=0, dilation=1, bias=False,
        )
        self.bn1 = nn.BatchNorm1d(bottleneck_channels)

        """         # 4) conv2: 3×3
        self.conv2 = spconv.SubMConv2d(
            bottleneck_channels, bottleneck_channels,
            kernel_size=3, stride=stride_3x3,
            padding=dilation, dilation=dilation,
            groups=num_groups, bias=False,
        ) """
        # 4) conv2: 3×3
        ConvClass = spconv.SparseConv2d if stride_3x3 > 1 else spconv.SubMConv2d
        self.conv2 = ConvClass(
            bottleneck_channels, bottleneck_channels,
            kernel_size=3, stride=stride_3x3,
            padding=dilation, dilation=dilation,
            groups=num_groups, bias=False,
        )
        self.bn2 = nn.BatchNorm1d(bottleneck_channels)

        # 5) conv3: 1×1
        self.conv3 = spconv.SubMConv2d(
            bottleneck_channels, out_channels,
            kernel_size=1, stride=1,
            padding=0, dilation=1, bias=False,
        )
        self.bn3 = nn.BatchNorm1d(out_channels)

        # 6) MSRA init (will be overwritten by your permuted .pth load)
        from fvcore.nn.weight_init import c2_msra_fill
        for conv in (self.conv1, self.conv2, self.conv3,
                     getattr(self, "shortcut_conv", None)):
            if conv is not None:
                c2_msra_fill(conv)

    @staticmethod
    def _dense_to_sparse(x):
        # identical logic to your stem helper
        B, C, H, W = x.shape
        nz_mask = x.reshape(B, C, -1).abs().sum(1) != 0
        idxs, feats = [], []
        for b in range(B):
            m = nz_mask[b].reshape(H, W)
            ys, xs = m.nonzero(as_tuple=True)
            if ys.numel() == 0:
                continue
            coords = torch.stack([
                torch.full_like(ys, b), ys, xs
            ], dim=1)
            idxs.append(coords)
            feats.append(x[b, :, ys, xs].t())
        idxs = torch.cat(idxs, 0).int()
        feats = torch.cat(feats, 0).contiguous()
        return spconv.SparseConvTensor(
            features=feats,
            indices=idxs,
            spatial_shape=(H, W),
            batch_size=B,
        )

    def forward(self, x_dense):
        # 1) dense → sparse
        st = self._dense_to_sparse(x_dense)

        # 2) shortcut path
        if self.shortcut_conv is not None:
            sc = self.shortcut_conv(st)
            sc = sc.replace_feature(self.shortcut_bn(sc.features))
        else:
            sc = st  # identity

        # 3) main path: conv1 → BN → ReLU
        out = self.conv1(st)
        out = out.replace_feature(self.bn1(out.features))
        out = out.replace_feature(F.relu(out.features))

        # 4) conv2 → BN → ReLU
        out = self.conv2(out)
        out = out.replace_feature(self.bn2(out.features))
        out = out.replace_feature(F.relu(out.features))

        # 5) conv3 → BN (no ReLU yet)
        out = self.conv3(out)
        out = out.replace_feature(self.bn3(out.features))

        # 6) add the shortcut features, then ReLU
        out = out.replace_feature(out.features + sc.features)
        out = out.replace_feature(F.relu(out.features))

        # 7) back to dense
        x_out = out.dense()      # (B, out_channels, H, W)
        return x_out 

class CustomBackboneWrapper(Backbone):
    def __init__(self, backbone):
        super().__init__()
        self.backbone = backbone

    def forward(self, x):
        return self.backbone(x)

        """     def forward(self, x):
        feats = self.backbone(x)   # this is the dict of “res2”, “res3”, “res4”, …
        for name, f in feats.items():
            # f.shape is (N, C, H, W)
            print(f"[backbone out] {name}: {tuple(f.shape)}==============================================================")
        return feats """

    def output_shape(self):
        shapes = self.backbone.output_shape()
        # Create new shapes dictionary while forcing specific stride values:
        new_shapes = {}
        if "res2" in shapes:
            # Force to original stride 4 (so that FPN computes log2(4)=2 -> "p2")
            new_shapes["res2"] = ShapeSpec(
                channels=shapes["res2"].channels,
                stride=1,  # force to 4
                height=shapes["res2"].height,
                width=shapes["res2"].width,
            )
        if "res3" in shapes:
            new_shapes["res3"] = ShapeSpec(
                channels=shapes["res3"].channels,
                stride=2,  # force to 8
                height=shapes["res3"].height,
                width=shapes["res3"].width,
            )
        if "res4" in shapes:
            new_shapes["res4"] = ShapeSpec(
                channels=shapes["res4"].channels,
                stride=4,  # force to 16
                height=shapes["res4"].height,
                width=shapes["res4"].width,
            )

        if "res5" in shapes:
            new_shapes["res5"] = ShapeSpec(
                channels=shapes["res5"].channels,
                stride=8,  # force to 32
                height=shapes["res5"].height,
                width=shapes["res5"].width,
            )
        return new_shapes


#from MIMOStem import MIMOStem
from SparseMIMOStem import SparseMIMOStem
@BACKBONE_REGISTRY.register()
class CustomResNetBackboneSparseA1(Backbone):
    def __init__(self, cfg, input_shape):
        super().__init__()  # Properly initialize the parent class
        # Build a standard ResNet backbone from the config with the given input_shape.
        _resnet_module.BottleneckBlock = SparseBottleneckBlock
        bottom_up = build_resnet_backbone(cfg, input_shape)
        bottom_up.stem = SparseMIMOStem(
            #in_channels=192,
            in_channels=input_shape.channels,
            out_channels=cfg.MODEL.RESNETS.STEM_OUT_CHANNELS,
            kernel_size=(3,13),
            dilation=(1,16),
            use_bn=True,  # or not, as you wish
            padding_ants=96,  # or something suiting your data
            stride=1,         # or 2, if you want downsampling
            norm=cfg.MODEL.RESNETS.NORM
        )
        bottom_up = CustomBackboneWrapper(bottom_up) 

        # Now, use the same configuration parameters:
        in_features = cfg.MODEL.FPN.IN_FEATURES
        out_channels = cfg.MODEL.FPN.OUT_CHANNELS

        # Manually build the FPN with top_block set to None to disable extra levels (p6)
        backbone = FPN(
            bottom_up=bottom_up,
            in_features=in_features,
            out_channels=out_channels,
            norm=cfg.MODEL.FPN.NORM,
            top_block=None,  # Disable the top block that would normally generate p6
            fuse_type=cfg.MODEL.FPN.FUSE_TYPE,
        )
        
        self.backbone = backbone

    def forward(self, x):
        return self.backbone(x)

    def output_shape(self):
        return self.backbone.output_shape()