The neurorient from zhantaochen

neurorient's People

Contributors

Watchers

neurorient's Issues

Model creation better be independent of external configs?

Just a general concern that currently models like ResNet2RotMat cannot be created without the _CONFIG. I still feel it is better to allow people create them by simply passing through args and kwargs.

I am going to inject _CONFIG through args and kwargs, feel free to let me know what you think @carbonscott

class ResNet2RotMat(nn.Module):
    def __init__(self, size=50, pretrained=False):
        super().__init__()
        resnet_type = f"resnet{size}"
        self.backbone = ImageEncoder(resnet_type, pretrained)

        # Create the adapter layer between backbone and bifpn...
        self.backbone_to_bifpn = nn.ModuleList([
            DepthwiseSeparableConv2d(in_channels  = in_channels,
                                     out_channels = _CONFIG.BIFPN.NUM_FEATURES,
                                     kernel_size  = 1,
                                     stride       = 1,
                                     padding      = 0)
            for _, in_channels in _CONFIG.BACKBONE.OUTPUT_CHANNELS.items()
        ])

        self.bifpn = BiFPN(num_blocks   = _CONFIG.BIFPN.NUM_BLOCKS,
                           num_features = _CONFIG.BIFPN.NUM_FEATURES,
                           num_levels   = _CONFIG.BIFPN.NUM_LEVELS)

        self.regressor_head = nn.Linear(_CONFIG.REGRESSOR_HEAD.IN_FEATURES, _CONFIG.REGRESSOR_HEAD.OUT_FEATURES)

Recommend Projects