KeyError: 'res4' using Swin backbone with detectron2 FCOS

[HiImCoder1818]

I am trying to get an FCOS detection model working with a Swin transformer backbone. I have tried to edit my yaml configuration to the best of my abilities by setting IN_FEATURES: ["p3", "p4", "p5", "p6"]. however I keep getting this error.

Traceback (most recent call last):
  File "/lambda/nfs/plane/ScoutingApp/detectors/artifacts/./fcos.py", line 39, in <module>
    model = build_model(cfg)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/detectron2/modeling/meta_arch/build.py", line 22, in build_model
    model = META_ARCH_REGISTRY.get(meta_arch)(cfg)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/detectron2/config/config.py", line 189, in wrapped
    explicit_args = _get_args_from_config(from_config_func, *args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/detectron2/config/config.py", line 245, in _get_args_from_config
    ret = from_config_func(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/detectron2/modeling/meta_arch/rcnn.py", line 76, in from_config
    "proposal_generator": build_proposal_generator(cfg, backbone.output_shape()),
  File "/home/ubuntu/.local/lib/python3.10/site-packages/detectron2/modeling/proposal_generator/build.py", line 24, in build_proposal_generator
    return PROPOSAL_GENERATOR_REGISTRY.get(name)(cfg, input_shape)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/detectron2/config/config.py", line 189, in wrapped
    explicit_args = _get_args_from_config(from_config_func, *args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/detectron2/config/config.py", line 245, in _get_args_from_config
    ret = from_config_func(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/detectron2/modeling/proposal_generator/rpn.py", line 280, in from_config
    ret["anchor_generator"] = build_anchor_generator(cfg, [input_shape[f] for f in in_features])
  File "/home/ubuntu/.local/lib/python3.10/site-packages/detectron2/modeling/proposal_generator/rpn.py", line 280, in <listcomp>
    ret["anchor_generator"] = build_anchor_generator(cfg, [input_shape[f] for f in in_features])
KeyError: 'res4'

My swin backbone is loaded like this:

from detectron2.config import get_cfg
from detectron2.modeling import build_model
from detectron2.modeling import BACKBONE_REGISTRY
import os

from swin import SwinBackbone, SwinArch

@BACKBONE_REGISTRY.register()
class SwinBackboneWrapper(SwinBackbone):
    def __init__(self, cfg, input_shape):
        swin_arch = SwinArch()
        super().__init__(swin_arch)

cfg = get_cfg()

path = "/lambda/nfs/plane/ScoutingApp/detectors/artifacts/config/fcos_R_50_FPN_1x.yaml"
cfg.set_new_allowed(True)
cfg.merge_from_file(path)
model = build_model(cfg)

And the SwinBackbone class is like:

class SwinBackbone(Backbone):
    def __init__(self, swin_arch):
        super().__init__()
        self.swin_arch = swin_arch
        self.strides = []
        self.in_channels = self.swin_arch.in_channels

    def forward(self, x):
        feature_maps = self.swin_arch(x)
        B, C, H, W = x.shape
        for fmap in feature_maps:
           self.strides.append(H // fmap.shape[2])

        return {f"p{i+3}": fmap for i, fmap in enumerate(feature_maps)}

    def output_shape(self):
        return {
            f"p{i+3}": ShapeSpec(channels=c, stride=s) for i, (c, s) in enumerate(zip(self.in_channels, self.strides))
        }

The swin arch is just a standard swin transformer that follows the paper. I later plan to train the model with a standard pytorch training loop and not with DefaultTrainer