RADAR ZeroGPU Space

.gitattributes

@@ -0,0 +1 @@
+*.pt filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,17 @@
+# Model weights and volumes are fetched at runtime, never committed.
+*.pth
+*.ptc
+*.nii
+*.nii.gz
+*.hdr
+*.img
+RADAR_infer_results_*.csv
+# Upstream scaffold, removed after vendoring.
+radar-upstream/
+damo-radar.zip
+# Python noise.
+__pycache__/
+*.pyc
+.venv/
+# Local knowledge-graph tooling; never pushed to the Space.
+graphify-out/

RADAR_inference/dynamic_network_architectures/architectures/resnet.py

@@ -0,0 +1,236 @@
+import torch
+from dynamic_network_architectures.building_blocks.residual_encoders import ResidualEncoder, BottleneckD, BasicBlockD
+from dynamic_network_architectures.building_blocks.helper import get_matching_pool_op, get_default_network_config
+from dynamic_network_architectures.building_blocks.simple_conv_blocks import ConvDropoutNormReLU
+from torch import nn
+
+_ResNet_CONFIGS = {
+    '18': {'features_per_stage': (64, 128, 256, 512), 'n_blocks_per_stage': (2, 2, 2, 2), 'strides': (1, 2, 2, 2),
+           'block': BasicBlockD, 'bottleneck_channels': None, 'disable_default_stem': True, 'stem_channels': None},
+    '34': {'features_per_stage': (64, 128, 256, 512), 'n_blocks_per_stage': (3, 4, 6, 3), 'strides': (1, 2, 2, 2),
+           'block': BasicBlockD, 'bottleneck_channels': None, 'disable_default_stem': True, 'stem_channels': None},
+    '50': {'features_per_stage': (64, 128, 256, 512), 'n_blocks_per_stage': (4, 6, 10, 5), 'strides': (1, 2, 2, 2),
+           'block': BasicBlockD, 'bottleneck_channels': None, 'disable_default_stem': True, 'stem_channels': None},
+    '152': {'features_per_stage': (64, 128, 256, 512), 'n_blocks_per_stage': (4, 13, 55, 4), 'strides': (1, 2, 2, 2),
+            'block': BasicBlockD, 'bottleneck_channels': None, 'disable_default_stem': True, 'stem_channels': None},
+    '50_bn': {'features_per_stage': (256, 512, 1024, 2048), 'n_blocks_per_stage': (3, 4, 6, 3), 'strides': (1, 2, 2, 2),
+              'block': BottleneckD, 'bottleneck_channels': (64, 128, 256, 512), 'disable_default_stem': True,
+              'stem_channels': 64},
+    '152_bn': {'features_per_stage': (256, 512, 1024, 2048), 'n_blocks_per_stage': (3, 8, 36, 3),
+               'strides': (1, 2, 2, 2),
+               'block': BottleneckD, 'bottleneck_channels': (64, 128, 256, 512), 'disable_default_stem': True,
+               'stem_channels': 64},
+    '18_cifar': {'features_per_stage': (64, 128, 256, 512), 'n_blocks_per_stage': (2, 2, 2, 2), 'strides': (1, 2, 2, 2),
+                 'block': BasicBlockD, 'bottleneck_channels': None, 'disable_default_stem': False,
+                 'stem_channels': None},
+    '34_cifar': {'features_per_stage': (64, 128, 256, 512), 'n_blocks_per_stage': (3, 4, 6, 3), 'strides': (1, 2, 2, 2),
+                 'block': BasicBlockD, 'bottleneck_channels': None, 'disable_default_stem': False,
+                 'stem_channels': None},
+    '50_cifar': {'features_per_stage': (64, 128, 256, 512), 'n_blocks_per_stage': (4, 6, 10, 5),
+                 'strides': (1, 2, 2, 2),
+                 'block': BasicBlockD, 'bottleneck_channels': None, 'disable_default_stem': False,
+                 'stem_channels': None},
+    '152_cifar': {'features_per_stage': (64, 128, 256, 512), 'n_blocks_per_stage': (4, 13, 55, 4),
+                  'strides': (1, 2, 2, 2),
+                  'block': BasicBlockD, 'bottleneck_channels': None, 'disable_default_stem': False,
+                  'stem_channels': None},
+    '50_cifar_bn': {'features_per_stage': (256, 512, 1024, 2048), 'n_blocks_per_stage': (3, 4, 6, 3),
+                    'strides': (1, 2, 2, 2),
+                    'block': BottleneckD, 'bottleneck_channels': (64, 128, 256, 512), 'disable_default_stem': False,
+                    'stem_channels': 64},
+    '152_cifar_bn': {'features_per_stage': (256, 512, 1024, 2048), 'n_blocks_per_stage': (3, 8, 36, 3),
+                     'strides': (1, 2, 2, 2),
+                     'block': BottleneckD, 'bottleneck_channels': (64, 128, 256, 512), 'disable_default_stem': False,
+                     'stem_channels': 64},
+}
+
+
+class ResNetD(nn.Module):
+    def __init__(self, n_classes: int, n_input_channel: int = 3, config='18', input_dimension=2,
+                 final_layer_dropout=0.0, stochastic_depth_p=0.0, squeeze_excitation=False,
+                 squeeze_excitation_rd_ratio=1./16):
+        """
+        Implements ResNetD (https://arxiv.org/pdf/1812.01187.pdf).
+        Args:
+            n_classes: Number of classes
+            n_input_channel: Number of input channels (e.g. 3 for RGB)
+            config: Configuration of the ResNet
+            input_dimension: Number of dimensions of the data (1, 2 or 3)
+            final_layer_dropout: Probability of dropout before the final classifier
+            stochastic_depth_p: Stochastic Depth probability
+            squeeze_excitation: Whether Squeeze and Excitation should be applied
+            squeeze_excitation_rd_ratio: Squeeze and Excitation Reduction Ratio
+        Returns:
+            ResNet Model
+        """
+        super().__init__()
+        self.input_channels = n_input_channel
+        self.cfg = _ResNet_CONFIGS[config]
+        self.ops = get_default_network_config(dimension=input_dimension)
+        self.final_layer_dropout_p = final_layer_dropout
+
+        if self.cfg['disable_default_stem']:
+            stem_features = self.cfg['stem_channels'] if self.cfg['stem_channels'] is not None else \
+            self.cfg['features_per_stage'][0]
+            self.stem = self._build_imagenet_stem_D(stem_features)
+            encoder_input_features = stem_features
+        else:
+            encoder_input_features = n_input_channel
+            self.stem = None
+
+        self.encoder = ResidualEncoder(encoder_input_features, n_stages=len(self.cfg['features_per_stage']),
+                                       features_per_stage=self.cfg['features_per_stage'], conv_op=self.ops['conv_op'],
+                                       kernel_sizes=3, strides=self.cfg['strides'],
+                                       n_blocks_per_stage=self.cfg['n_blocks_per_stage'], conv_bias=False,
+                                       norm_op=self.ops['norm_op'], norm_op_kwargs=None, dropout_op=None,
+                                       dropout_op_kwargs=None, nonlin=nn.ReLU,
+                                       nonlin_kwargs={'inplace': True}, block=self.cfg['block'],
+                                       bottleneck_channels=self.cfg['bottleneck_channels'], return_skips=False,
+                                       disable_default_stem=self.cfg['disable_default_stem'],
+                                       stem_channels=self.cfg['stem_channels'],
+                                       stochastic_depth_p=stochastic_depth_p,
+                                       squeeze_excitation=squeeze_excitation,
+                                       squeeze_excitation_reduction_ratio=squeeze_excitation_rd_ratio)
+
+        self.gap = get_matching_pool_op(conv_op=self.ops['conv_op'], adaptive=True, pool_type='avg')(1)
+        self.classifier = nn.Linear(self.cfg['features_per_stage'][-1], n_classes, True)
+        self.final_layer_dropout = self.ops['dropout_op'](p=self.final_layer_dropout_p)
+
+    def forward(self, x):
+        if self.stem is not None:
+            x = self.stem(x)
+        x = self.encoder(x)
+        x = self.gap(x)
+        x = self.final_layer_dropout(x).squeeze()
+
+        return self.classifier(x)
+
+    def _build_imagenet_stem_D(self, stem_features):
+        """
+        https://arxiv.org/pdf/1812.01187.pdf
+
+        use 3 3x3(x3) convs instead of one 7x7. Stride is located in first conv.
+
+        Fig2 b) describes this
+        :return:
+        """
+        c1 = ConvDropoutNormReLU(self.ops['conv_op'], self.input_channels, stem_features, 3, 2, False,
+                                 self.ops['norm_op'], None, None, None, nn.ReLU, {'inplace': True})
+        c2 = ConvDropoutNormReLU(self.ops['conv_op'], stem_features, stem_features, 3, 1, False,
+                                 self.ops['norm_op'], None, None, None, nn.ReLU, {'inplace': True})
+        c3 = ConvDropoutNormReLU(self.ops['conv_op'], stem_features, stem_features, 3, 1, False,
+                                 self.ops['norm_op'], None, None, None, nn.ReLU, {'inplace': True})
+        pl = get_matching_pool_op(conv_op=self.ops['conv_op'], adaptive=False, pool_type='max')(2)
+        stem = nn.Sequential(c1, c2, c3, pl)
+        return stem
+
+
+class ResNet18_CIFAR(ResNetD):
+    def __init__(self, n_classes: int, n_input_channels: int = 3, input_dimension: int = 2,
+                 final_layer_dropout: float = 0.0, stochastic_depth_p: float = 0.0, squeeze_excitation: bool = False,
+                 squeeze_excitation_rd_ratio: float = 1./16):
+        super().__init__(n_classes, n_input_channels, config='18_cifar', input_dimension=input_dimension,
+                         final_layer_dropout=final_layer_dropout, stochastic_depth_p=stochastic_depth_p,
+                         squeeze_excitation=squeeze_excitation, squeeze_excitation_rd_ratio=squeeze_excitation_rd_ratio)
+
+class ResNet34_CIFAR(ResNetD):
+    def __init__(self, n_classes: int, n_input_channels: int = 3, input_dimension: int = 2,
+                 final_layer_dropout: float = 0.0, stochastic_depth_p: float = 0.0, squeeze_excitation: bool = False,
+                 squeeze_excitation_rd_ratio: float = 1./16):
+        super().__init__(n_classes, n_input_channels, config='34_cifar', input_dimension=input_dimension,
+                         final_layer_dropout=final_layer_dropout, stochastic_depth_p=stochastic_depth_p,
+                         squeeze_excitation=squeeze_excitation, squeeze_excitation_rd_ratio=squeeze_excitation_rd_ratio)
+
+class ResNet50_CIFAR(ResNetD):
+    def __init__(self, n_classes: int, n_input_channels: int = 3, input_dimension: int = 2,
+                 final_layer_dropout: float = 0.0, stochastic_depth_p: float = 0.0, squeeze_excitation: bool = False,
+                 squeeze_excitation_rd_ratio: float = 1./16):
+        super().__init__(n_classes, n_input_channels, config='50_cifar', input_dimension=input_dimension,
+                         final_layer_dropout=final_layer_dropout, stochastic_depth_p=stochastic_depth_p,
+                         squeeze_excitation=squeeze_excitation, squeeze_excitation_rd_ratio=squeeze_excitation_rd_ratio)
+
+class ResNet152_CIFAR(ResNetD):
+    def __init__(self, n_classes: int, n_input_channels: int = 3, input_dimension: int = 2,
+                 final_layer_dropout: float = 0.0, stochastic_depth_p: float = 0.0, squeeze_excitation: bool = False,
+                 squeeze_excitation_rd_ratio: float = 1./16):
+        super().__init__(n_classes, n_input_channels, config='152_cifar', input_dimension=input_dimension,
+                         final_layer_dropout=final_layer_dropout, stochastic_depth_p=stochastic_depth_p,
+                         squeeze_excitation=squeeze_excitation, squeeze_excitation_rd_ratio=squeeze_excitation_rd_ratio)
+
+class ResNet50bn_CIFAR(ResNetD):
+    def __init__(self, n_classes: int, n_input_channels: int = 3, input_dimension: int = 2,
+                 final_layer_dropout: float = 0.0, stochastic_depth_p: float = 0.0, squeeze_excitation: bool = False,
+                 squeeze_excitation_rd_ratio: float = 1./16):
+        super().__init__(n_classes, n_input_channels, config='50_cifar_bn', input_dimension=input_dimension,
+                         final_layer_dropout=final_layer_dropout, stochastic_depth_p=stochastic_depth_p,
+                         squeeze_excitation=squeeze_excitation, squeeze_excitation_rd_ratio=squeeze_excitation_rd_ratio)
+
+class ResNet152bn_CIFAR(ResNetD):
+    def __init__(self, n_classes: int, n_input_channels: int = 3, input_dimension: int = 2,
+                 final_layer_dropout: float = 0.0, stochastic_depth_p: float = 0.0, squeeze_excitation: bool = False,
+                 squeeze_excitation_rd_ratio: float = 1./16):
+        super().__init__(n_classes, n_input_channels, config='152_cifar_bn', input_dimension=input_dimension,
+                         final_layer_dropout=final_layer_dropout, stochastic_depth_p=stochastic_depth_p,
+                         squeeze_excitation=squeeze_excitation, squeeze_excitation_rd_ratio=squeeze_excitation_rd_ratio)
+
+class ResNet18(ResNetD):
+    def __init__(self, n_classes: int, n_input_channels: int = 3, input_dimension: int = 2,
+                 final_layer_dropout: float = 0.0, stochastic_depth_p: float = 0.0, squeeze_excitation: bool = False,
+                 squeeze_excitation_rd_ratio: float = 1./16):
+        super().__init__(n_classes, n_input_channels, config='18', input_dimension=input_dimension,
+                         final_layer_dropout=final_layer_dropout, stochastic_depth_p=stochastic_depth_p,
+                         squeeze_excitation=squeeze_excitation, squeeze_excitation_rd_ratio=squeeze_excitation_rd_ratio)
+
+class ResNet34(ResNetD):
+    def __init__(self, n_classes: int, n_input_channels: int = 3, input_dimension: int = 2,
+                 final_layer_dropout: float = 0.0, stochastic_depth_p: float = 0.0, squeeze_excitation: bool = False,
+                 squeeze_excitation_rd_ratio: float = 1./16):
+        super().__init__(n_classes, n_input_channels, config='34', input_dimension=input_dimension,
+                         final_layer_dropout=final_layer_dropout, stochastic_depth_p=stochastic_depth_p,
+                         squeeze_excitation=squeeze_excitation, squeeze_excitation_rd_ratio=squeeze_excitation_rd_ratio)
+
+class ResNet50(ResNetD):
+    def __init__(self, n_classes: int, n_input_channels: int = 3, input_dimension: int = 2,
+                 final_layer_dropout: float = 0.0, stochastic_depth_p: float = 0.0, squeeze_excitation: bool = False,
+                 squeeze_excitation_rd_ratio: float = 1./16):
+        super().__init__(n_classes, n_input_channels, config='50', input_dimension=input_dimension,
+                         final_layer_dropout=final_layer_dropout, stochastic_depth_p=stochastic_depth_p,
+                         squeeze_excitation=squeeze_excitation, squeeze_excitation_rd_ratio=squeeze_excitation_rd_ratio)
+
+class ResNet152(ResNetD):
+    def __init__(self, n_classes: int, n_input_channels: int = 3, input_dimension: int = 2,
+                 final_layer_dropout: float = 0.0, stochastic_depth_p: float = 0.0, squeeze_excitation: bool = False,
+                 squeeze_excitation_rd_ratio: float = 1./16):
+        super().__init__(n_classes, n_input_channels, config='152', input_dimension=input_dimension,
+                         final_layer_dropout=final_layer_dropout, stochastic_depth_p=stochastic_depth_p,
+                         squeeze_excitation=squeeze_excitation, squeeze_excitation_rd_ratio=squeeze_excitation_rd_ratio)
+
+class ResNet50bn(ResNetD):
+    def __init__(self, n_classes: int, n_input_channels: int = 3, input_dimension: int = 2,
+                 final_layer_dropout: float = 0.0, stochastic_depth_p: float = 0.0, squeeze_excitation: bool = False,
+                 squeeze_excitation_rd_ratio: float = 1./16):
+        super().__init__(n_classes, n_input_channels, config='50_bn', input_dimension=input_dimension,
+                         final_layer_dropout=final_layer_dropout, stochastic_depth_p=stochastic_depth_p,
+                         squeeze_excitation=squeeze_excitation, squeeze_excitation_rd_ratio=squeeze_excitation_rd_ratio)
+
+class ResNet152bn(ResNetD):
+    def __init__(self, n_classes: int, n_input_channels: int = 3, input_dimension: int = 2,
+                 final_layer_dropout: float = 0.0, stochastic_depth_p: float = 0.0, squeeze_excitation: bool = False,
+                 squeeze_excitation_rd_ratio: float = 1./16):
+        super().__init__(n_classes, n_input_channels, config='152_bn', input_dimension=input_dimension,
+                         final_layer_dropout=final_layer_dropout, stochastic_depth_p=stochastic_depth_p,
+                         squeeze_excitation=squeeze_excitation, squeeze_excitation_rd_ratio=squeeze_excitation_rd_ratio)
+
+
+if __name__ == '__main__':
+    data = torch.rand((1, 3, 224, 224))
+
+    model = ResNet50bn(10, 3)
+    import hiddenlayer as hl
+
+    g = hl.build_graph(model, data,
+                       transforms=None)
+    g.save("network_architecture.pdf")
+    del g
+
+    #print(model.compute_conv_feature_map_size((32, 32)))

RADAR_inference/dynamic_network_architectures/architectures/resnet_vl.py

@@ -0,0 +1,225 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+import math
+from functools import partial
+
+
+def conv3x3x3(in_planes, out_planes, stride=1, dilation=1):
+    # 3x3x3 convolution with padding
+    return nn.Conv3d(
+        in_planes,
+        out_planes,
+        kernel_size=3,
+        dilation=dilation,
+        stride=stride,
+        padding=dilation,
+        bias=False)
+
+
+def downsample_basic_block(x, planes, stride, no_cuda=False):
+    out = F.avg_pool3d(x, kernel_size=1, stride=stride)
+    zero_pads = torch.Tensor(
+        out.size(0), planes - out.size(1), out.size(2), out.size(3),
+        out.size(4)).zero_()
+    if not no_cuda:
+        # if isinstance(out.data, torch.cuda.FloatTensor):
+        zero_pads = zero_pads.cuda()
+
+    out = Variable(torch.cat([out.data, zero_pads], dim=1))
+
+    return out
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3x3(inplanes, planes, stride=stride, dilation=dilation)
+        self.bn1 = nn.BatchNorm3d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3x3(planes, planes, dilation=dilation)
+        self.bn2 = nn.BatchNorm3d(planes)
+        self.downsample = downsample
+        self.stride = stride
+        self.dilation = dilation
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv3d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm3d(planes)
+        self.conv2 = nn.Conv3d(
+            planes, planes, kernel_size=3, stride=stride, dilation=dilation, padding=dilation, bias=False)
+        self.bn2 = nn.BatchNorm3d(planes)
+        self.conv3 = nn.Conv3d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm3d(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+        self.dilation = dilation
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+class ResNet(nn.Module):
+
+    def __init__(self,
+                 block,
+                 layers,
+                 shortcut_type='B',
+                 no_cuda = False):
+        self.inplanes = 64
+        self.no_cuda = no_cuda
+        super(ResNet, self).__init__()
+        self.conv1 = nn.Conv3d(
+            1,
+            64,
+            kernel_size=7,
+            stride=(1, 2, 2),
+            padding=(3, 3, 3),
+            bias=False)
+        
+        self.bn1 = nn.BatchNorm3d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool3d(kernel_size=(3, 3, 3), stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0], shortcut_type)
+        self.layer2 = self._make_layer(
+            block, 128, layers[1], shortcut_type, stride=2)
+        self.layer3 = self._make_layer(
+            block, 256, layers[2], shortcut_type, stride=2, dilation=2)
+        self.layer4 = self._make_layer(
+            block, 512, layers[3], shortcut_type, stride=2, dilation=4)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv3d):
+                m.weight = nn.init.kaiming_normal(m.weight, mode='fan_out')
+            elif isinstance(m, nn.BatchNorm3d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, shortcut_type, stride=1, dilation=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            if shortcut_type == 'A':
+                downsample = partial(
+                    downsample_basic_block,
+                    planes=planes * block.expansion,
+                    stride=stride,
+                    no_cuda=self.no_cuda)
+            else:
+                downsample = nn.Sequential(
+                    nn.Conv3d(
+                        self.inplanes,
+                        planes * block.expansion,
+                        kernel_size=1,
+                        stride=stride,
+                        bias=False), nn.BatchNorm3d(planes * block.expansion))
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride=stride, dilation=dilation, downsample=downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes, dilation=dilation))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x6 = self.relu(x)
+        x5 = self.maxpool(x6)
+        x4 = self.layer1(x5)
+        x3 = self.layer2(x4)
+        x2 = self.layer3(x3)
+        x1 = self.layer4(x2)
+
+        return [x6, x4, x3, x2, x1]
+
+def resnet10(**kwargs):
+    """Constructs a ResNet-18 model.
+    """
+    model = ResNet(BasicBlock, [1, 1, 1, 1], **kwargs)
+    return model
+
+def resnet18(**kwargs):
+    """Constructs a ResNet-18 model.
+    """
+    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
+    return model
+
+
+def resnet34(**kwargs):
+    """Constructs a ResNet-34 model.
+    """
+    model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
+    return model
+
+
+def resnet50(**kwargs):
+    """Constructs a ResNet-50 model.
+    """
+    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
+    return model
+
+
+def resnet101(**kwargs):
+    """Constructs a ResNet-101 model.
+    """
+    model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
+    return model
+
+
+def resnet152(**kwargs):
+    """Constructs a ResNet-101 model.
+    """
+    model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
+    return model
+
+
+def resnet200(**kwargs):
+    """Constructs a ResNet-101 model.
+    """
+    model = ResNet(Bottleneck, [3, 24, 36, 3], **kwargs)
+    return model

RADAR_inference/dynamic_network_architectures/architectures/unet.py

@@ -0,0 +1,218 @@
+from typing import Union, Type, List, Tuple
+
+import torch
+from dynamic_network_architectures.building_blocks.helper import convert_conv_op_to_dim
+from dynamic_network_architectures.building_blocks.plain_conv_encoder import PlainConvEncoder
+from dynamic_network_architectures.building_blocks.residual import BasicBlockD, BottleneckD
+from dynamic_network_architectures.building_blocks.residual_encoders import ResidualEncoder
+from dynamic_network_architectures.building_blocks.unet_decoder import UNetDecoder
+from dynamic_network_architectures.building_blocks.unet_residual_decoder import UNetResDecoder
+from dynamic_network_architectures.initialization.weight_init import InitWeights_He
+from dynamic_network_architectures.initialization.weight_init import init_last_bn_before_add_to_0
+from torch import nn
+from torch.nn.modules.conv import _ConvNd
+from torch.nn.modules.dropout import _DropoutNd
+
+
+class PlainConvUNet(nn.Module):
+    def __init__(self,
+                 input_channels: int,
+                 n_stages: int,
+                 features_per_stage: Union[int, List[int], Tuple[int, ...]],
+                 conv_op: Type[_ConvNd],
+                 kernel_sizes: Union[int, List[int], Tuple[int, ...]],
+                 strides: Union[int, List[int], Tuple[int, ...]],
+                 n_conv_per_stage: Union[int, List[int], Tuple[int, ...]],
+                 num_classes: int,
+                 n_conv_per_stage_decoder: Union[int, Tuple[int, ...], List[int]],
+                 conv_bias: bool = False,
+                 norm_op: Union[None, Type[nn.Module]] = None,
+                 norm_op_kwargs: dict = None,
+                 dropout_op: Union[None, Type[_DropoutNd]] = None,
+                 dropout_op_kwargs: dict = None,
+                 nonlin: Union[None, Type[torch.nn.Module]] = None,
+                 nonlin_kwargs: dict = None,
+                 deep_supervision: bool = False,
+                 nonlin_first: bool = False
+                 ):
+        """
+        nonlin_first: if True you get conv -> nonlin -> norm. Else it's conv -> norm -> nonlin
+        """
+        super().__init__()
+        if isinstance(n_conv_per_stage, int):
+            n_conv_per_stage = [n_conv_per_stage] * n_stages
+        if isinstance(n_conv_per_stage_decoder, int):
+            n_conv_per_stage_decoder = [n_conv_per_stage_decoder] * (n_stages - 1)
+        assert len(n_conv_per_stage) == n_stages, "n_conv_per_stage must have as many entries as we have " \
+                                                  f"resolution stages. here: {n_stages}. " \
+                                                  f"n_conv_per_stage: {n_conv_per_stage}"
+        assert len(n_conv_per_stage_decoder) == (n_stages - 1), "n_conv_per_stage_decoder must have one less entries " \
+                                                                f"as we have resolution stages. here: {n_stages} " \
+                                                                f"stages, so it should have {n_stages - 1} entries. " \
+                                                                f"n_conv_per_stage_decoder: {n_conv_per_stage_decoder}"
+        self.encoder = PlainConvEncoder(input_channels, n_stages, features_per_stage, conv_op, kernel_sizes, strides,
+                                        n_conv_per_stage, conv_bias, norm_op, norm_op_kwargs, dropout_op,
+                                        dropout_op_kwargs, nonlin, nonlin_kwargs, return_skips=True,
+                                        nonlin_first=nonlin_first)
+        self.decoder = UNetDecoder(self.encoder, num_classes, n_conv_per_stage_decoder, deep_supervision,
+                                   nonlin_first=nonlin_first)
+
+    def forward(self, x):
+        skips = self.encoder(x)
+        return self.decoder(skips)
+
+    def compute_conv_feature_map_size(self, input_size):
+        assert len(input_size) == convert_conv_op_to_dim(self.encoder.conv_op), "just give the image size without color/feature channels or " \
+                                                            "batch channel. Do not give input_size=(b, c, x, y(, z)). " \
+                                                            "Give input_size=(x, y(, z))!"
+        return self.encoder.compute_conv_feature_map_size(input_size) + self.decoder.compute_conv_feature_map_size(input_size)
+
+    @staticmethod
+    def initialize(module):
+        InitWeights_He(1e-2)(module)
+
+
+class ResidualEncoderUNet(nn.Module):
+    def __init__(self,
+                 input_channels: int,
+                 n_stages: int,
+                 features_per_stage: Union[int, List[int], Tuple[int, ...]],
+                 conv_op: Type[_ConvNd],
+                 kernel_sizes: Union[int, List[int], Tuple[int, ...]],
+                 strides: Union[int, List[int], Tuple[int, ...]],
+                 n_blocks_per_stage: Union[int, List[int], Tuple[int, ...]],
+                 num_classes: int,
+                 n_conv_per_stage_decoder: Union[int, Tuple[int, ...], List[int]],
+                 conv_bias: bool = False,
+                 norm_op: Union[None, Type[nn.Module]] = None,
+                 norm_op_kwargs: dict = None,
+                 dropout_op: Union[None, Type[_DropoutNd]] = None,
+                 dropout_op_kwargs: dict = None,
+                 nonlin: Union[None, Type[torch.nn.Module]] = None,
+                 nonlin_kwargs: dict = None,
+                 deep_supervision: bool = False,
+                 block: Union[Type[BasicBlockD], Type[BottleneckD]] = BasicBlockD,
+                 bottleneck_channels: Union[int, List[int], Tuple[int, ...]] = None,
+                 stem_channels: int = None
+                 ):
+        super().__init__()
+        if isinstance(n_blocks_per_stage, int):
+            n_blocks_per_stage = [n_blocks_per_stage] * n_stages
+        if isinstance(n_conv_per_stage_decoder, int):
+            n_conv_per_stage_decoder = [n_conv_per_stage_decoder] * (n_stages - 1)
+        assert len(n_blocks_per_stage) == n_stages, "n_blocks_per_stage must have as many entries as we have " \
+                                                  f"resolution stages. here: {n_stages}. " \
+                                                  f"n_blocks_per_stage: {n_blocks_per_stage}"
+        assert len(n_conv_per_stage_decoder) == (n_stages - 1), "n_conv_per_stage_decoder must have one less entries " \
+                                                                f"as we have resolution stages. here: {n_stages} " \
+                                                                f"stages, so it should have {n_stages - 1} entries. " \
+                                                                f"n_conv_per_stage_decoder: {n_conv_per_stage_decoder}"
+        self.encoder = ResidualEncoder(input_channels, n_stages, features_per_stage, conv_op, kernel_sizes, strides,
+                                       n_blocks_per_stage, conv_bias, norm_op, norm_op_kwargs, dropout_op,
+                                       dropout_op_kwargs, nonlin, nonlin_kwargs, block, bottleneck_channels,
+                                       return_skips=True, disable_default_stem=False, stem_channels=stem_channels)
+        self.decoder = UNetDecoder(self.encoder, num_classes, n_conv_per_stage_decoder, deep_supervision)
+
+    def forward(self, x):
+        skips = self.encoder(x)
+        return self.decoder(skips)
+
+    def compute_conv_feature_map_size(self, input_size):
+        assert len(input_size) == convert_conv_op_to_dim(self.encoder.conv_op), "just give the image size without color/feature channels or " \
+                                                                                "batch channel. Do not give input_size=(b, c, x, y(, z)). " \
+                                                                                "Give input_size=(x, y(, z))!"
+        return self.encoder.compute_conv_feature_map_size(input_size) + self.decoder.compute_conv_feature_map_size(input_size)
+
+    @staticmethod
+    def initialize(module):
+        InitWeights_He(1e-2)(module)
+        init_last_bn_before_add_to_0(module)
+
+
+class ResidualUNet(nn.Module):
+    def __init__(self,
+                 input_channels: int,
+                 n_stages: int,
+                 features_per_stage: Union[int, List[int], Tuple[int, ...]],
+                 conv_op: Type[_ConvNd],
+                 kernel_sizes: Union[int, List[int], Tuple[int, ...]],
+                 strides: Union[int, List[int], Tuple[int, ...]],
+                 n_blocks_per_stage: Union[int, List[int], Tuple[int, ...]],
+                 num_classes: int,
+                 n_conv_per_stage_decoder: Union[int, Tuple[int, ...], List[int]],
+                 conv_bias: bool = False,
+                 norm_op: Union[None, Type[nn.Module]] = None,
+                 norm_op_kwargs: dict = None,
+                 dropout_op: Union[None, Type[_DropoutNd]] = None,
+                 dropout_op_kwargs: dict = None,
+                 nonlin: Union[None, Type[torch.nn.Module]] = None,
+                 nonlin_kwargs: dict = None,
+                 deep_supervision: bool = False,
+                 block: Union[Type[BasicBlockD], Type[BottleneckD]] = BasicBlockD,
+                 bottleneck_channels: Union[int, List[int], Tuple[int, ...]] = None,
+                 stem_channels: int = None
+                 ):
+        super().__init__()
+        if isinstance(n_blocks_per_stage, int):
+            n_blocks_per_stage = [n_blocks_per_stage] * n_stages
+        if isinstance(n_conv_per_stage_decoder, int):
+            n_conv_per_stage_decoder = [n_conv_per_stage_decoder] * (n_stages - 1)
+        assert len(n_blocks_per_stage) == n_stages, "n_blocks_per_stage must have as many entries as we have " \
+                                                  f"resolution stages. here: {n_stages}. " \
+                                                  f"n_blocks_per_stage: {n_blocks_per_stage}"
+        assert len(n_conv_per_stage_decoder) == (n_stages - 1), "n_conv_per_stage_decoder must have one less entries " \
+                                                                f"as we have resolution stages. here: {n_stages} " \
+                                                                f"stages, so it should have {n_stages - 1} entries. " \
+                                                                f"n_conv_per_stage_decoder: {n_conv_per_stage_decoder}"
+        self.encoder = ResidualEncoder(input_channels, n_stages, features_per_stage, conv_op, kernel_sizes, strides,
+                                       n_blocks_per_stage, conv_bias, norm_op, norm_op_kwargs, dropout_op,
+                                       dropout_op_kwargs, nonlin, nonlin_kwargs, block, bottleneck_channels,
+                                       return_skips=True, disable_default_stem=False, stem_channels=stem_channels)
+        self.decoder = UNetResDecoder(self.encoder, num_classes, n_conv_per_stage_decoder, deep_supervision)
+
+    def forward(self, x):
+        skips = self.encoder(x)
+        return self.decoder(skips)
+
+    def compute_conv_feature_map_size(self, input_size):
+        assert len(input_size) == convert_conv_op_to_dim(self.encoder.conv_op), "just give the image size without color/feature channels or " \
+                                                                                "batch channel. Do not give input_size=(b, c, x, y(, z)). " \
+                                                                                "Give input_size=(x, y(, z))!"
+        return self.encoder.compute_conv_feature_map_size(input_size) + self.decoder.compute_conv_feature_map_size(input_size)
+
+    @staticmethod
+    def initialize(module):
+        InitWeights_He(1e-2)(module)
+        init_last_bn_before_add_to_0(module)
+
+
+if __name__ == '__main__':
+    data = torch.rand((1, 4, 128, 128, 128))
+
+    model = PlainConvUNet(4, 6, (32, 64, 125, 256, 320, 320), nn.Conv3d, 3, (1, 2, 2, 2, 2, 2), (2, 2, 2, 2, 2, 2), 4,
+                                (2, 2, 2, 2, 2), False, nn.BatchNorm3d, None, None, None, nn.ReLU, deep_supervision=True)
+
+    if False:
+        import hiddenlayer as hl
+
+        g = hl.build_graph(model, data,
+                           transforms=None)
+        g.save("network_architecture.pdf")
+        del g
+
+    print(model.compute_conv_feature_map_size(data.shape[2:]))
+
+    data = torch.rand((1, 4, 512, 512))
+
+    model = PlainConvUNet(4, 8, (32, 64, 125, 256, 512, 512, 512, 512), nn.Conv2d, 3, (1, 2, 2, 2, 2, 2, 2, 2), (2, 2, 2, 2, 2, 2, 2, 2), 4,
+                                (2, 2, 2, 2, 2, 2, 2), False, nn.BatchNorm2d, None, None, None, nn.ReLU, deep_supervision=True)
+
+    if False:
+        import hiddenlayer as hl
+
+        g = hl.build_graph(model, data,
+                           transforms=None)
+        g.save("network_architecture.pdf")
+        del g
+
+    print(model.compute_conv_feature_map_size(data.shape[2:]))

RADAR_inference/dynamic_network_architectures/architectures/unet_lightdecoder.py

@@ -0,0 +1,218 @@
+from typing import Union, Type, List, Tuple
+
+import torch
+from dynamic_network_architectures.building_blocks.helper import convert_conv_op_to_dim
+from dynamic_network_architectures.building_blocks.plain_conv_encoder import PlainConvEncoder
+from dynamic_network_architectures.building_blocks.residual import BasicBlockD, BottleneckD
+from dynamic_network_architectures.building_blocks.residual_encoders import ResidualEncoder
+from dynamic_network_architectures.building_blocks.unet_decoder_light import UNetDecoder
+from dynamic_network_architectures.building_blocks.unet_residual_decoder import UNetResDecoder
+from dynamic_network_architectures.initialization.weight_init import InitWeights_He
+from dynamic_network_architectures.initialization.weight_init import init_last_bn_before_add_to_0
+from torch import nn
+from torch.nn.modules.conv import _ConvNd
+from torch.nn.modules.dropout import _DropoutNd
+
+
+class PlainConvUNetLightD(nn.Module):
+    def __init__(self,
+                 input_channels: int,
+                 n_stages: int,
+                 features_per_stage: Union[int, List[int], Tuple[int, ...]],
+                 conv_op: Type[_ConvNd],
+                 kernel_sizes: Union[int, List[int], Tuple[int, ...]],
+                 strides: Union[int, List[int], Tuple[int, ...]],
+                 n_conv_per_stage: Union[int, List[int], Tuple[int, ...]],
+                 num_classes: int,
+                 n_conv_per_stage_decoder: Union[int, Tuple[int, ...], List[int]],
+                 conv_bias: bool = False,
+                 norm_op: Union[None, Type[nn.Module]] = None,
+                 norm_op_kwargs: dict = None,
+                 dropout_op: Union[None, Type[_DropoutNd]] = None,
+                 dropout_op_kwargs: dict = None,
+                 nonlin: Union[None, Type[torch.nn.Module]] = None,
+                 nonlin_kwargs: dict = None,
+                 deep_supervision: bool = False,
+                 nonlin_first: bool = False
+                 ):
+        """
+        nonlin_first: if True you get conv -> nonlin -> norm. Else it's conv -> norm -> nonlin
+        """
+        super().__init__()
+        if isinstance(n_conv_per_stage, int):
+            n_conv_per_stage = [n_conv_per_stage] * n_stages
+        if isinstance(n_conv_per_stage_decoder, int):
+            n_conv_per_stage_decoder = [n_conv_per_stage_decoder] * (n_stages - 1)
+        assert len(n_conv_per_stage) == n_stages, "n_conv_per_stage must have as many entries as we have " \
+                                                  f"resolution stages. here: {n_stages}. " \
+                                                  f"n_conv_per_stage: {n_conv_per_stage}"
+        assert len(n_conv_per_stage_decoder) == (n_stages - 1), "n_conv_per_stage_decoder must have one less entries " \
+                                                                f"as we have resolution stages. here: {n_stages} " \
+                                                                f"stages, so it should have {n_stages - 1} entries. " \
+                                                                f"n_conv_per_stage_decoder: {n_conv_per_stage_decoder}"
+        self.encoder = PlainConvEncoder(input_channels, n_stages, features_per_stage, conv_op, kernel_sizes, strides,
+                                        n_conv_per_stage, conv_bias, norm_op, norm_op_kwargs, dropout_op,
+                                        dropout_op_kwargs, nonlin, nonlin_kwargs, return_skips=True,
+                                        nonlin_first=nonlin_first)
+        self.decoder = UNetDecoder(self.encoder, num_classes, n_conv_per_stage_decoder, deep_supervision,
+                                   nonlin_first=nonlin_first)
+
+    def forward(self, x):
+        skips = self.encoder(x)
+        return skips, self.decoder(skips)
+
+    def compute_conv_feature_map_size(self, input_size):
+        assert len(input_size) == convert_conv_op_to_dim(self.encoder.conv_op), "just give the image size without color/feature channels or " \
+                                                            "batch channel. Do not give input_size=(b, c, x, y(, z)). " \
+                                                            "Give input_size=(x, y(, z))!"
+        return self.encoder.compute_conv_feature_map_size(input_size) + self.decoder.compute_conv_feature_map_size(input_size)
+
+    @staticmethod
+    def initialize(module):
+        InitWeights_He(1e-2)(module)
+
+
+class ResidualEncoderUNet(nn.Module):
+    def __init__(self,
+                 input_channels: int,
+                 n_stages: int,
+                 features_per_stage: Union[int, List[int], Tuple[int, ...]],
+                 conv_op: Type[_ConvNd],
+                 kernel_sizes: Union[int, List[int], Tuple[int, ...]],
+                 strides: Union[int, List[int], Tuple[int, ...]],
+                 n_blocks_per_stage: Union[int, List[int], Tuple[int, ...]],
+                 num_classes: int,
+                 n_conv_per_stage_decoder: Union[int, Tuple[int, ...], List[int]],
+                 conv_bias: bool = False,
+                 norm_op: Union[None, Type[nn.Module]] = None,
+                 norm_op_kwargs: dict = None,
+                 dropout_op: Union[None, Type[_DropoutNd]] = None,
+                 dropout_op_kwargs: dict = None,
+                 nonlin: Union[None, Type[torch.nn.Module]] = None,
+                 nonlin_kwargs: dict = None,
+                 deep_supervision: bool = False,
+                 block: Union[Type[BasicBlockD], Type[BottleneckD]] = BasicBlockD,
+                 bottleneck_channels: Union[int, List[int], Tuple[int, ...]] = None,
+                 stem_channels: int = None
+                 ):
+        super().__init__()
+        if isinstance(n_blocks_per_stage, int):
+            n_blocks_per_stage = [n_blocks_per_stage] * n_stages
+        if isinstance(n_conv_per_stage_decoder, int):
+            n_conv_per_stage_decoder = [n_conv_per_stage_decoder] * (n_stages - 1)
+        assert len(n_blocks_per_stage) == n_stages, "n_blocks_per_stage must have as many entries as we have " \
+                                                  f"resolution stages. here: {n_stages}. " \
+                                                  f"n_blocks_per_stage: {n_blocks_per_stage}"
+        assert len(n_conv_per_stage_decoder) == (n_stages - 1), "n_conv_per_stage_decoder must have one less entries " \
+                                                                f"as we have resolution stages. here: {n_stages} " \
+                                                                f"stages, so it should have {n_stages - 1} entries. " \
+                                                                f"n_conv_per_stage_decoder: {n_conv_per_stage_decoder}"
+        self.encoder = ResidualEncoder(input_channels, n_stages, features_per_stage, conv_op, kernel_sizes, strides,
+                                       n_blocks_per_stage, conv_bias, norm_op, norm_op_kwargs, dropout_op,
+                                       dropout_op_kwargs, nonlin, nonlin_kwargs, block, bottleneck_channels,
+                                       return_skips=True, disable_default_stem=False, stem_channels=stem_channels)
+        self.decoder = UNetDecoder(self.encoder, num_classes, n_conv_per_stage_decoder, deep_supervision)
+
+    def forward(self, x):
+        skips = self.encoder(x)
+        return self.decoder(skips)
+
+    def compute_conv_feature_map_size(self, input_size):
+        assert len(input_size) == convert_conv_op_to_dim(self.encoder.conv_op), "just give the image size without color/feature channels or " \
+                                                                                "batch channel. Do not give input_size=(b, c, x, y(, z)). " \
+                                                                                "Give input_size=(x, y(, z))!"
+        return self.encoder.compute_conv_feature_map_size(input_size) + self.decoder.compute_conv_feature_map_size(input_size)
+
+    @staticmethod
+    def initialize(module):
+        InitWeights_He(1e-2)(module)
+        init_last_bn_before_add_to_0(module)
+
+
+class ResidualUNet(nn.Module):
+    def __init__(self,
+                 input_channels: int,
+                 n_stages: int,
+                 features_per_stage: Union[int, List[int], Tuple[int, ...]],
+                 conv_op: Type[_ConvNd],
+                 kernel_sizes: Union[int, List[int], Tuple[int, ...]],
+                 strides: Union[int, List[int], Tuple[int, ...]],
+                 n_blocks_per_stage: Union[int, List[int], Tuple[int, ...]],
+                 num_classes: int,
+                 n_conv_per_stage_decoder: Union[int, Tuple[int, ...], List[int]],
+                 conv_bias: bool = False,
+                 norm_op: Union[None, Type[nn.Module]] = None,
+                 norm_op_kwargs: dict = None,
+                 dropout_op: Union[None, Type[_DropoutNd]] = None,
+                 dropout_op_kwargs: dict = None,
+                 nonlin: Union[None, Type[torch.nn.Module]] = None,
+                 nonlin_kwargs: dict = None,
+                 deep_supervision: bool = False,
+                 block: Union[Type[BasicBlockD], Type[BottleneckD]] = BasicBlockD,
+                 bottleneck_channels: Union[int, List[int], Tuple[int, ...]] = None,
+                 stem_channels: int = None
+                 ):
+        super().__init__()
+        if isinstance(n_blocks_per_stage, int):
+            n_blocks_per_stage = [n_blocks_per_stage] * n_stages
+        if isinstance(n_conv_per_stage_decoder, int):
+            n_conv_per_stage_decoder = [n_conv_per_stage_decoder] * (n_stages - 1)
+        assert len(n_blocks_per_stage) == n_stages, "n_blocks_per_stage must have as many entries as we have " \
+                                                  f"resolution stages. here: {n_stages}. " \
+                                                  f"n_blocks_per_stage: {n_blocks_per_stage}"
+        assert len(n_conv_per_stage_decoder) == (n_stages - 1), "n_conv_per_stage_decoder must have one less entries " \
+                                                                f"as we have resolution stages. here: {n_stages} " \
+                                                                f"stages, so it should have {n_stages - 1} entries. " \
+                                                                f"n_conv_per_stage_decoder: {n_conv_per_stage_decoder}"
+        self.encoder = ResidualEncoder(input_channels, n_stages, features_per_stage, conv_op, kernel_sizes, strides,
+                                       n_blocks_per_stage, conv_bias, norm_op, norm_op_kwargs, dropout_op,
+                                       dropout_op_kwargs, nonlin, nonlin_kwargs, block, bottleneck_channels,
+                                       return_skips=True, disable_default_stem=False, stem_channels=stem_channels)
+        self.decoder = UNetResDecoder(self.encoder, num_classes, n_conv_per_stage_decoder, deep_supervision)
+
+    def forward(self, x):
+        skips = self.encoder(x)
+        return self.decoder(skips)
+
+    def compute_conv_feature_map_size(self, input_size):
+        assert len(input_size) == convert_conv_op_to_dim(self.encoder.conv_op), "just give the image size without color/feature channels or " \
+                                                                                "batch channel. Do not give input_size=(b, c, x, y(, z)). " \
+                                                                                "Give input_size=(x, y(, z))!"
+        return self.encoder.compute_conv_feature_map_size(input_size) + self.decoder.compute_conv_feature_map_size(input_size)
+
+    @staticmethod
+    def initialize(module):
+        InitWeights_He(1e-2)(module)
+        init_last_bn_before_add_to_0(module)
+
+
+if __name__ == '__main__':
+    data = torch.rand((1, 4, 128, 128, 128))
+
+    model = PlainConvUNet(4, 6, (32, 64, 125, 256, 320, 320), nn.Conv3d, 3, (1, 2, 2, 2, 2, 2), (2, 2, 2, 2, 2, 2), 4,
+                                (2, 2, 2, 2, 2), False, nn.BatchNorm3d, None, None, None, nn.ReLU, deep_supervision=True)
+
+    if False:
+        import hiddenlayer as hl
+
+        g = hl.build_graph(model, data,
+                           transforms=None)
+        g.save("network_architecture.pdf")
+        del g
+
+    print(model.compute_conv_feature_map_size(data.shape[2:]))
+
+    data = torch.rand((1, 4, 512, 512))
+
+    model = PlainConvUNet(4, 8, (32, 64, 125, 256, 512, 512, 512, 512), nn.Conv2d, 3, (1, 2, 2, 2, 2, 2, 2, 2), (2, 2, 2, 2, 2, 2, 2, 2), 4,
+                                (2, 2, 2, 2, 2, 2, 2), False, nn.BatchNorm2d, None, None, None, nn.ReLU, deep_supervision=True)
+
+    if False:
+        import hiddenlayer as hl
+
+        g = hl.build_graph(model, data,
+                           transforms=None)
+        g.save("network_architecture.pdf")
+        del g
+
+    print(model.compute_conv_feature_map_size(data.shape[2:]))

RADAR_inference/dynamic_network_architectures/architectures/vgg.py

@@ -0,0 +1,85 @@
+import torch
+from torch import nn
+
+from dynamic_network_architectures.building_blocks.plain_conv_encoder import PlainConvEncoder
+from dynamic_network_architectures.building_blocks.helper import get_matching_pool_op, get_default_network_config
+
+_VGG_CONFIGS = {
+    '16': {'features_per_stage': (64, 128, 256, 512, 512, 512), 'n_conv_per_stage': (2, 2, 2, 3, 3, 3),
+           'strides': (1, 2, 2, 2, 2, 2)},
+    '19': {'features_per_stage': (64, 128, 256, 512, 512, 512), 'n_conv_per_stage': (2, 2, 3, 3, 4, 4),
+           'strides': (1, 2, 2, 2, 2, 2)},
+    '16_cifar': {'features_per_stage': (64, 128, 256, 512), 'n_conv_per_stage': (2, 3, 5, 5), 'strides': (1, 2, 2, 2)},
+    '19_cifar': {'features_per_stage': (64, 128, 256, 512), 'n_conv_per_stage': (3, 4, 5, 6), 'strides': (1, 2, 2, 2)},
+}
+
+_VGG_OPS = {
+    1: {'conv_op': nn.Conv1d, 'norm_op': nn.BatchNorm1d},
+    2: {'conv_op': nn.Conv2d, 'norm_op': nn.BatchNorm2d},
+    3: {'conv_op': nn.Conv3d, 'norm_op': nn.BatchNorm3d},
+}
+
+
+class VGG(nn.Module):
+    def __init__(self, n_classes: int, n_input_channel: int = 3, config='16', input_dimension=2):
+        """
+        This is not 1:1 VGG because it does not have the bloated fully connected layers at the end. Since these were
+        counted towards the XX layers as well, we increase the number of convolutional layers so that we have the
+        desired number of conv layers in total
+
+        We also use batchnorm
+        """
+        super().__init__()
+        cfg = _VGG_CONFIGS[config]
+        ops = get_default_network_config(dimension=input_dimension)
+        self.encoder = PlainConvEncoder(
+            n_input_channel, n_stages=len(cfg['features_per_stage']), features_per_stage=cfg['features_per_stage'],
+            conv_op=ops['conv_op'],
+            kernel_sizes=3, strides=cfg['strides'], n_conv_per_stage=cfg['n_conv_per_stage'], conv_bias=False,
+            norm_op=ops['norm_op'], norm_op_kwargs=None, dropout_op=None, dropout_op_kwargs=None, nonlin=nn.ReLU,
+            nonlin_kwargs={'inplace': True}, return_skips=False
+        )
+        self.gap = get_matching_pool_op(conv_op=ops['conv_op'], adaptive=True, pool_type='avg')(1)
+        self.classifier = nn.Linear(cfg['features_per_stage'][-1], n_classes, True)
+
+    def forward(self, x):
+        x = self.encoder(x)
+        x = self.gap(x).squeeze()
+        return self.classifier(x)
+
+    def compute_conv_feature_map_size(self, input_size):
+        return self.encoder.compute_conv_feature_map_size(input_size)
+
+
+class VGG16(VGG):
+    def __init__(self, n_classes: int, n_input_channel: int = 3, input_dimension: int = 2):
+        super().__init__(n_classes, n_input_channel, config='16', input_dimension=input_dimension)
+
+
+class VGG19(VGG):
+    def __init__(self, n_classes: int, n_input_channel: int = 3, input_dimension: int = 2):
+        super().__init__(n_classes, n_input_channel, config='19', input_dimension=input_dimension)
+
+
+class VGG16_cifar(VGG):
+    def __init__(self, n_classes: int, n_input_channel: int = 3, input_dimension: int = 2):
+        super().__init__(n_classes, n_input_channel, config='16_cifar', input_dimension=input_dimension)
+
+
+class VGG19_cifar(VGG):
+    def __init__(self, n_classes: int, n_input_channel: int = 3, input_dimension: int = 2):
+        super().__init__(n_classes, n_input_channel, config='19_cifar', input_dimension=input_dimension)
+
+
+if __name__ == '__main__':
+    data = torch.rand((1, 3, 32, 32))
+
+    model = VGG19_cifar(10, 3)
+    import hiddenlayer as hl
+
+    g = hl.build_graph(model, data,
+                       transforms=None)
+    g.save("network_architecture.pdf")
+    del g
+
+    print(model.compute_conv_feature_map_size((32, 32)))

RADAR_inference/dynamic_network_architectures/building_blocks/helper.py

@@ -0,0 +1,242 @@
+from typing import Type
+import numpy as np
+import torch.nn
+from torch import nn
+from torch.nn.modules.batchnorm import _BatchNorm
+from torch.nn.modules.conv import _ConvNd, _ConvTransposeNd
+from torch.nn.modules.dropout import _DropoutNd
+from torch.nn.modules.instancenorm import _InstanceNorm
+
+
+def convert_dim_to_conv_op(dimension: int) -> Type[_ConvNd]:
+    """
+    :param dimension: 1, 2 or 3
+    :return: conv Class of corresponding dimension
+    """
+    if dimension == 1:
+        return nn.Conv1d
+    elif dimension == 2:
+        return nn.Conv2d
+    elif dimension == 3:
+        return nn.Conv3d
+    else:
+        raise ValueError("Unknown dimension. Only 1, 2 and 3 are supported")
+
+
+def convert_conv_op_to_dim(conv_op: Type[_ConvNd]) -> int:
+    """
+    :param conv_op: conv class
+    :return: dimension: 1, 2 or 3
+    """
+    if conv_op == nn.Conv1d:
+        return 1
+    elif conv_op == nn.Conv2d:
+        return 2
+    elif conv_op == nn.Conv3d:
+        return 3
+    else:
+        raise ValueError("Unknown dimension. Only 1d 2d and 3d conv are supported. got %s" % str(conv_op))
+
+
+def get_matching_pool_op(conv_op: Type[_ConvNd] = None,
+                         dimension: int = None,
+                         adaptive=False,
+                         pool_type: str = 'avg') -> Type[torch.nn.Module]:
+    """
+    You MUST set EITHER conv_op OR dimension. Do not set both!
+    :param conv_op:
+    :param dimension:
+    :param adaptive:
+    :param pool_type: either 'avg' or 'max'
+    :return:
+    """
+    assert not ((conv_op is not None) and (dimension is not None)), \
+        "You MUST set EITHER conv_op OR dimension. Do not set both!"
+    assert pool_type in ['avg', 'max'], 'pool_type must be either avg or max'
+    if conv_op is not None:
+        dimension = convert_conv_op_to_dim(conv_op)
+    assert dimension in [1, 2, 3], 'Dimension must be 1, 2 or 3'
+
+    if conv_op is not None:
+        dimension = convert_conv_op_to_dim(conv_op)
+
+    if dimension == 1:
+        if pool_type == 'avg':
+            if adaptive:
+                return nn.AdaptiveAvgPool1d
+            else:
+                return nn.AvgPool1d
+        elif pool_type == 'max':
+            if adaptive:
+                return nn.AdaptiveMaxPool1d
+            else:
+                return nn.MaxPool1d
+    elif dimension == 2:
+        if pool_type == 'avg':
+            if adaptive:
+                return nn.AdaptiveAvgPool2d
+            else:
+                return nn.AvgPool2d
+        elif pool_type == 'max':
+            if adaptive:
+                return nn.AdaptiveMaxPool2d
+            else:
+                return nn.MaxPool2d
+    elif dimension == 3:
+        if pool_type == 'avg':
+            if adaptive:
+                return nn.AdaptiveAvgPool3d
+            else:
+                return nn.AvgPool3d
+        elif pool_type == 'max':
+            if adaptive:
+                return nn.AdaptiveMaxPool3d
+            else:
+                return nn.MaxPool3d
+
+
+def get_matching_instancenorm(conv_op: Type[_ConvNd] = None, dimension: int = None) -> Type[_InstanceNorm]:
+    """
+    You MUST set EITHER conv_op OR dimension. Do not set both!
+
+    :param conv_op:
+    :param dimension:
+    :return:
+    """
+    assert not ((conv_op is not None) and (dimension is not None)), \
+        "You MUST set EITHER conv_op OR dimension. Do not set both!"
+    if conv_op is not None:
+        dimension = convert_conv_op_to_dim(conv_op)
+    if dimension is not None:
+        assert dimension in [1, 2, 3], 'Dimension must be 1, 2 or 3'
+    if dimension == 1:
+        return nn.InstanceNorm1d
+    elif dimension == 2:
+        return nn.InstanceNorm2d
+    elif dimension == 3:
+        return nn.InstanceNorm3d
+
+
+def get_matching_convtransp(conv_op: Type[_ConvNd] = None, dimension: int = None) -> Type[_ConvTransposeNd]:
+    """
+    You MUST set EITHER conv_op OR dimension. Do not set both!
+
+    :param conv_op:
+    :param dimension:
+    :return:
+    """
+    assert not ((conv_op is not None) and (dimension is not None)), \
+        "You MUST set EITHER conv_op OR dimension. Do not set both!"
+    if conv_op is not None:
+        dimension = convert_conv_op_to_dim(conv_op)
+    assert dimension in [1, 2, 3], 'Dimension must be 1, 2 or 3'
+    if dimension == 1:
+        return nn.ConvTranspose1d
+    elif dimension == 2:
+        return nn.ConvTranspose2d
+    elif dimension == 3:
+        return nn.ConvTranspose3d
+
+
+def get_matching_batchnorm(conv_op: Type[_ConvNd] = None, dimension: int = None) -> Type[_BatchNorm]:
+    """
+    You MUST set EITHER conv_op OR dimension. Do not set both!
+
+    :param conv_op:
+    :param dimension:
+    :return:
+    """
+    assert not ((conv_op is not None) and (dimension is not None)), \
+        "You MUST set EITHER conv_op OR dimension. Do not set both!"
+    if conv_op is not None:
+        dimension = convert_conv_op_to_dim(conv_op)
+    assert dimension in [1, 2, 3], 'Dimension must be 1, 2 or 3'
+    if dimension == 1:
+        return nn.BatchNorm1d
+    elif dimension == 2:
+        return nn.BatchNorm2d
+    elif dimension == 3:
+        return nn.BatchNorm3d
+
+
+def get_matching_dropout(conv_op: Type[_ConvNd] = None, dimension: int = None) -> Type[_DropoutNd]:
+    """
+    You MUST set EITHER conv_op OR dimension. Do not set both!
+
+    :param conv_op:
+    :param dimension:
+    :return:
+    """
+    assert not ((conv_op is not None) and (dimension is not None)), \
+        "You MUST set EITHER conv_op OR dimension. Do not set both!"
+    assert dimension in [1, 2, 3], 'Dimension must be 1, 2 or 3'
+    if dimension == 1:
+        return nn.Dropout
+    elif dimension == 2:
+        return nn.Dropout2d
+    elif dimension == 3:
+        return nn.Dropout3d
+
+
+def maybe_convert_scalar_to_list(conv_op, scalar):
+    """
+    useful for converting, for example, kernel_size=3 to [3, 3, 3] in case of nn.Conv3d
+    :param conv_op:
+    :param scalar:
+    :return:
+    """
+    if not isinstance(scalar, (tuple, list, np.ndarray)):
+        if conv_op == nn.Conv2d:
+            return [scalar] * 2
+        elif conv_op == nn.Conv3d:
+            return [scalar] * 3
+        elif conv_op == nn.Conv1d:
+            return [scalar] * 1
+        else:
+            raise RuntimeError("Invalid conv op: %s" % str(conv_op))
+    else:
+        return scalar
+
+
+def get_default_network_config(dimension: int = 2,
+                               nonlin: str = "ReLU",
+                               norm_type: str = "bn") -> dict:
+    """
+    Use this to get a standard configuration. A network configuration looks like this:
+
+    config = {'conv_op': torch.nn.modules.conv.Conv2d,
+              'dropout_op': torch.nn.modules.dropout.Dropout2d,
+              'norm_op': torch.nn.modules.batchnorm.BatchNorm2d,
+              'norm_op_kwargs': {'eps': 1e-05, 'affine': True},
+              'nonlin': torch.nn.modules.activation.ReLU,
+              'nonlin_kwargs': {'inplace': True}}
+
+    There is no need to use get_default_network_config. You can create your own. Network configs are a convenient way of
+    setting dimensionality, normalization and nonlinearity.
+
+    :param dimension: integer denoting the dimension of the data. 1, 2 and 3 are accepted
+    :param nonlin: string (ReLU or LeakyReLU)
+    :param norm_type: string (bn=batch norm, in=instance norm)
+    torch.nn.Module
+    :return: dict
+    """
+    config = {}
+    config['conv_op'] = convert_dim_to_conv_op(dimension)
+    config['dropout_op'] = get_matching_dropout(dimension=dimension)
+    if norm_type == "bn":
+        config['norm_op'] = get_matching_batchnorm(dimension=dimension)
+    elif norm_type == "in":
+        config['norm_op'] = get_matching_instancenorm(dimension=dimension)
+
+    config['norm_op_kwargs'] = None # this will use defaults
+
+    if nonlin == "LeakyReLU":
+        config['nonlin'] = nn.LeakyReLU
+        config['nonlin_kwargs'] = {'negative_slope': 1e-2, 'inplace': True}
+    elif nonlin == "ReLU":
+        config['nonlin'] = nn.ReLU
+        config['nonlin_kwargs'] = {'inplace': True}
+    else:
+        raise NotImplementedError('Unknown nonlin %s. Only "LeakyReLU" and "ReLU" are supported for now' % nonlin)
+
+    return config

RADAR_inference/dynamic_network_architectures/building_blocks/plain_conv_encoder.py

@@ -0,0 +1,105 @@
+import torch
+from torch import nn
+import numpy as np
+from typing import Union, Type, List, Tuple
+
+from torch.nn.modules.conv import _ConvNd
+from torch.nn.modules.dropout import _DropoutNd
+from dynamic_network_architectures.building_blocks.simple_conv_blocks import StackedConvBlocks
+from dynamic_network_architectures.building_blocks.helper import maybe_convert_scalar_to_list, get_matching_pool_op
+
+
+class PlainConvEncoder(nn.Module):
+    def __init__(self,
+                 input_channels: int,
+                 n_stages: int,
+                 features_per_stage: Union[int, List[int], Tuple[int, ...]],
+                 conv_op: Type[_ConvNd],
+                 kernel_sizes: Union[int, List[int], Tuple[int, ...]],
+                 strides: Union[int, List[int], Tuple[int, ...]],
+                 n_conv_per_stage: Union[int, List[int], Tuple[int, ...]],
+                 conv_bias: bool = False,
+                 norm_op: Union[None, Type[nn.Module]] = None,
+                 norm_op_kwargs: dict = None,
+                 dropout_op: Union[None, Type[_DropoutNd]] = None,
+                 dropout_op_kwargs: dict = None,
+                 nonlin: Union[None, Type[torch.nn.Module]] = None,
+                 nonlin_kwargs: dict = None,
+                 return_skips: bool = False,
+                 nonlin_first: bool = False,
+                 pool: str = 'conv'
+                 ):
+
+        super().__init__()
+        if isinstance(kernel_sizes, int):
+            kernel_sizes = [kernel_sizes] * n_stages
+        if isinstance(features_per_stage, int):
+            features_per_stage = [features_per_stage] * n_stages
+        if isinstance(n_conv_per_stage, int):
+            n_conv_per_stage = [n_conv_per_stage] * n_stages
+        if isinstance(strides, int):
+            strides = [strides] * n_stages
+        assert len(kernel_sizes) == n_stages, "kernel_sizes must have as many entries as we have resolution stages (n_stages)"
+        assert len(n_conv_per_stage) == n_stages, "n_conv_per_stage must have as many entries as we have resolution stages (n_stages)"
+        assert len(features_per_stage) == n_stages, "features_per_stage must have as many entries as we have resolution stages (n_stages)"
+        assert len(strides) == n_stages, "strides must have as many entries as we have resolution stages (n_stages). " \
+                                             "Important: first entry is recommended to be 1, else we run strided conv drectly on the input"
+
+        stages = []
+        for s in range(n_stages):
+            stage_modules = []
+            if pool == 'max' or pool == 'avg':
+                if (isinstance(strides[s], int) and strides[s] != 1) or \
+                        isinstance(strides[s], (tuple, list)) and any([i != 1 for i in strides[s]]):
+                    stage_modules.append(get_matching_pool_op(conv_op, pool_type=pool)(kernel_size=strides[s], stride=strides[s]))
+                conv_stride = 1
+            elif pool == 'conv':
+                conv_stride = strides[s]
+            else:
+                raise RuntimeError()
+            stage_modules.append(StackedConvBlocks(
+                n_conv_per_stage[s], conv_op, input_channels, features_per_stage[s], kernel_sizes[s], conv_stride,
+                conv_bias, norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs, nonlin, nonlin_kwargs, nonlin_first
+            ))
+            stages.append(nn.Sequential(*stage_modules))
+            input_channels = features_per_stage[s]
+
+        self.stages = nn.Sequential(*stages)
+        self.output_channels = features_per_stage
+        self.strides = [maybe_convert_scalar_to_list(conv_op, i) for i in strides]
+        self.return_skips = return_skips
+
+        # we store some things that a potential decoder needs
+        self.conv_op = conv_op
+        self.norm_op = norm_op
+        self.norm_op_kwargs = norm_op_kwargs
+        self.nonlin = nonlin
+        self.nonlin_kwargs = nonlin_kwargs
+        self.dropout_op = dropout_op
+        self.dropout_op_kwargs = dropout_op_kwargs
+        self.conv_bias = conv_bias
+        self.kernel_sizes = kernel_sizes
+
+    def forward(self, x):
+        ret = []
+        for s in self.stages:
+            x = s(x)
+            ret.append(x)
+        if self.return_skips:
+            return ret
+        else:
+            return ret[-1]
+
+    def compute_conv_feature_map_size(self, input_size):
+        output = np.int64(0)
+        for s in range(len(self.stages)):
+            if isinstance(self.stages[s], nn.Sequential):
+                for sq in self.stages[s]:
+                    if hasattr(sq, 'compute_conv_feature_map_size'):
+                        output += self.stages[s][-1].compute_conv_feature_map_size(input_size)
+            else:
+                output += self.stages[s].compute_conv_feature_map_size(input_size)
+            input_size = [i // j for i, j in zip(input_size, self.strides[s])]
+        return output
+
+

RADAR_inference/dynamic_network_architectures/building_blocks/regularization.py

@@ -0,0 +1,86 @@
+from torch import nn
+
+
+def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by_keep: bool = True):
+    """
+    This function is taken from the timm package (https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py).
+
+    Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+    """
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+
+class DropPath(nn.Module):
+    """
+    This class is taken from the timm package (https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py).
+
+    Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
+
+
+class SqueezeExcite(nn.Module):
+    """
+    This class is taken from the timm package (https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/squeeze_excite.py)
+    and slightly modified so that the convolution type can be adapted.
+
+    SE Module as defined in original SE-Nets with a few additions
+    Additions include:
+        * divisor can be specified to keep channels % div == 0 (default: 8)
+        * reduction channels can be specified directly by arg (if rd_channels is set)
+        * reduction channels can be specified by float rd_ratio (default: 1/16)
+        * global max pooling can be added to the squeeze aggregation
+        * customizable activation, normalization, and gate layer
+    """
+    def __init__(
+            self, channels, conv_op, rd_ratio=1. / 16, rd_channels=None, rd_divisor=8, add_maxpool=False,
+            act_layer=nn.ReLU, norm_layer=None, gate_layer=nn.Sigmoid):
+        super(SqueezeExcite, self).__init__()
+        self.add_maxpool = add_maxpool
+        if not rd_channels:
+            rd_channels = make_divisible(channels * rd_ratio, rd_divisor, round_limit=0.)
+        self.fc1 = conv_op(channels, rd_channels, kernel_size=1, bias=True)
+        self.bn = norm_layer(rd_channels) if norm_layer else nn.Identity()
+        self.act = act_layer(inplace=True)
+        self.fc2 = conv_op(rd_channels, channels, kernel_size=1, bias=True)
+        self.gate = gate_layer()
+
+    def forward(self, x):
+        x_se = x.mean((2, 3), keepdim=True)
+        if self.add_maxpool:
+            # experimental codepath, may remove or change
+            x_se = 0.5 * x_se + 0.5 * x.amax((2, 3), keepdim=True)
+        x_se = self.fc1(x_se)
+        x_se = self.act(self.bn(x_se))
+        x_se = self.fc2(x_se)
+        return x * self.gate(x_se)
+
+
+def make_divisible(v, divisor=8, min_value=None, round_limit=.9):
+    """
+    This function is taken from the timm package (https://github.com/rwightman/pytorch-image-models/blob/b7cb8d0337b3e7b50516849805ddb9be5fc11644/timm/models/layers/helpers.py#L25)
+    """
+    min_value = min_value or divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < round_limit * v:
+        new_v += divisor
+    return new_v

RADAR_inference/dynamic_network_architectures/building_blocks/residual.py

@@ -0,0 +1,371 @@
+from typing import Tuple, List, Union, Type
+import torch.nn
+from torch import nn
+from torch.nn.modules.conv import _ConvNd
+from torch.nn.modules.dropout import _DropoutNd
+
+from dynamic_network_architectures.building_blocks.helper import maybe_convert_scalar_to_list, get_matching_pool_op
+from dynamic_network_architectures.building_blocks.simple_conv_blocks import ConvDropoutNormReLU
+from dynamic_network_architectures.building_blocks.regularization import DropPath, SqueezeExcite
+import numpy as np
+
+
+class BasicBlockD(nn.Module):
+    def __init__(self,
+                 conv_op: Type[_ConvNd],
+                 input_channels: int,
+                 output_channels: int,
+                 kernel_size: Union[int, List[int], Tuple[int, ...]],
+                 stride: Union[int, List[int], Tuple[int, ...]],
+                 conv_bias: bool = False,
+                 norm_op: Union[None, Type[nn.Module]] = None,
+                 norm_op_kwargs: dict = None,
+                 dropout_op: Union[None, Type[_DropoutNd]] = None,
+                 dropout_op_kwargs: dict = None,
+                 nonlin: Union[None, Type[torch.nn.Module]] = None,
+                 nonlin_kwargs: dict = None,
+                 stochastic_depth_p: float = 0.0,
+                 squeeze_excitation: bool = False,
+                 squeeze_excitation_reduction_ratio: float = 1. / 16,
+                 # todo wideresnet?
+                 ):
+        """
+        This implementation follows ResNet-D:
+
+        He, Tong, et al. "Bag of tricks for image classification with convolutional neural networks."
+        Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2019.
+
+        The skip has an avgpool (if needed) followed by 1x1 conv instead of just a strided 1x1 conv
+
+        :param conv_op:
+        :param input_channels:
+        :param output_channels:
+        :param kernel_size: refers only to convs in feature extraction path, not to 1x1x1 conv in skip
+        :param stride: only applies to first conv (and skip). Second conv always has stride 1
+        :param conv_bias:
+        :param norm_op:
+        :param norm_op_kwargs:
+        :param dropout_op: only the first conv can have dropout. The second never has
+        :param dropout_op_kwargs:
+        :param nonlin:
+        :param nonlin_kwargs:
+        :param stochastic_depth_p:
+        :param squeeze_excitation:
+        :param squeeze_excitation_reduction_ratio:
+        """
+        super().__init__()
+        self.input_channels = input_channels
+        self.output_channels = output_channels
+        stride = maybe_convert_scalar_to_list(conv_op, stride)
+        self.stride = stride
+
+        kernel_size = maybe_convert_scalar_to_list(conv_op, kernel_size)
+
+        if norm_op_kwargs is None:
+            norm_op_kwargs = {}
+        if nonlin_kwargs is None:
+            nonlin_kwargs = {}
+
+        self.conv1 = ConvDropoutNormReLU(conv_op, input_channels, output_channels, kernel_size, stride, conv_bias,
+                                         norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs, nonlin, nonlin_kwargs)
+        self.conv2 = ConvDropoutNormReLU(conv_op, output_channels, output_channels, kernel_size, 1, conv_bias, norm_op,
+                                         norm_op_kwargs, None, None, None, None)
+
+        self.nonlin2 = nonlin(**nonlin_kwargs) if nonlin is not None else lambda x: x
+
+        # Stochastic Depth
+        self.apply_stochastic_depth = False if stochastic_depth_p == 0.0 else True
+        if self.apply_stochastic_depth:
+            self.drop_path = DropPath(drop_prob=stochastic_depth_p)
+
+        # Squeeze Excitation
+        self.apply_se = squeeze_excitation
+        if self.apply_se:
+            self.squeeze_excitation = SqueezeExcite(self.output_channels, conv_op,
+                                                    rd_ratio=squeeze_excitation_reduction_ratio, rd_divisor=8)
+
+        has_stride = (isinstance(stride, int) and stride != 1) or any([i != 1 for i in stride])
+        requires_projection = (input_channels != output_channels)
+
+        if has_stride or requires_projection:
+            ops = []
+            if has_stride:
+                ops.append(get_matching_pool_op(conv_op=conv_op, adaptive=False, pool_type='avg')(stride, stride))
+            if requires_projection:
+                ops.append(
+                    ConvDropoutNormReLU(conv_op, input_channels, output_channels, 1, 1, False, norm_op,
+                                        norm_op_kwargs, None, None, None, None
+                                        )
+                )
+            self.skip = nn.Sequential(*ops)
+        else:
+            self.skip = lambda x: x
+
+    def forward(self, x):
+        residual = self.skip(x)
+        out = self.conv2(self.conv1(x))
+        if self.apply_stochastic_depth:
+            out = self.drop_path(out)
+        if self.apply_se:
+            out = self.squeeze_excitation(out)
+        out += residual
+        return self.nonlin2(out)
+
+    def compute_conv_feature_map_size(self, input_size):
+        assert len(input_size) == len(self.stride), "just give the image size without color/feature channels or " \
+                                                    "batch channel. Do not give input_size=(b, c, x, y(, z)). " \
+                                                    "Give input_size=(x, y(, z))!"
+        size_after_stride = [i // j for i, j in zip(input_size, self.stride)]
+        # conv1
+        output_size_conv1 = np.prod([self.output_channels, *size_after_stride], dtype=np.int64)
+        # conv2
+        output_size_conv2 = np.prod([self.output_channels, *size_after_stride], dtype=np.int64)
+        # skip conv (if applicable)
+        if (self.input_channels != self.output_channels) or any([i != j for i, j in zip(input_size, size_after_stride)]):
+            assert isinstance(self.skip, nn.Sequential)
+            output_size_skip = np.prod([self.output_channels, *size_after_stride], dtype=np.int64)
+        else:
+            assert not isinstance(self.skip, nn.Sequential)
+            output_size_skip = 0
+        return output_size_conv1 + output_size_conv2 + output_size_skip
+
+
+class BottleneckD(nn.Module):
+    def __init__(self,
+                 conv_op: Type[_ConvNd],
+                 input_channels: int,
+                 bottleneck_channels: int,
+                 output_channels: int,
+                 kernel_size: Union[int, List[int], Tuple[int, ...]],
+                 stride: Union[int, List[int], Tuple[int, ...]],
+                 conv_bias: bool = False,
+                 norm_op: Union[None, Type[nn.Module]] = None,
+                 norm_op_kwargs: dict = None,
+                 dropout_op: Union[None, Type[_DropoutNd]] = None,
+                 dropout_op_kwargs: dict = None,
+                 nonlin: Union[None, Type[torch.nn.Module]] = None,
+                 nonlin_kwargs: dict = None,
+                 stochastic_depth_p: float = 0.0,
+                 squeeze_excitation: bool = False,
+                 squeeze_excitation_reduction_ratio: float = 1. / 16
+                 ):
+        """
+        This implementation follows ResNet-D:
+
+        He, Tong, et al. "Bag of tricks for image classification with convolutional neural networks."
+        Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2019.
+
+        The stride sits in the 3x3 conv instead of the 1x1 conv!
+        The skip has an avgpool (if needed) followed by 1x1 conv instead of just a strided 1x1 conv
+
+        :param conv_op:
+        :param input_channels:
+        :param output_channels:
+        :param kernel_size: only affects the conv in the middle (typically 3x3). The other convs remain 1x1
+        :param stride: only applies to the conv in the middle (and skip). Note that this deviates from the canonical
+        ResNet implementation where the stride is applied to the first 1x1 conv. (This implementation follows ResNet-D)
+        :param conv_bias:
+        :param norm_op:
+        :param norm_op_kwargs:
+        :param dropout_op: only the second (kernel_size) conv can have dropout. The first and last conv (1x1(x1)) never have it
+        :param dropout_op_kwargs:
+        :param nonlin:
+        :param nonlin_kwargs:
+        :param stochastic_depth_p:
+        :param squeeze_excitation:
+        :param squeeze_excitation_reduction_ratio:
+        """
+        super().__init__()
+        self.input_channels = input_channels
+        self.output_channels = output_channels
+        self.bottleneck_channels = bottleneck_channels
+        stride = maybe_convert_scalar_to_list(conv_op, stride)
+        self.stride = stride
+
+        kernel_size = maybe_convert_scalar_to_list(conv_op, kernel_size)
+        if norm_op_kwargs is None:
+            norm_op_kwargs = {}
+        if nonlin_kwargs is None:
+            nonlin_kwargs = {}
+
+        self.conv1 = ConvDropoutNormReLU(conv_op, input_channels, bottleneck_channels, 1, 1, conv_bias,
+                                         norm_op, norm_op_kwargs, None, None, nonlin, nonlin_kwargs)
+        self.conv2 = ConvDropoutNormReLU(conv_op, bottleneck_channels, bottleneck_channels, kernel_size, stride,
+                                         conv_bias,
+                                         norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs, nonlin, nonlin_kwargs)
+        self.conv3 = ConvDropoutNormReLU(conv_op, bottleneck_channels, output_channels, 1, 1, conv_bias, norm_op,
+                                         norm_op_kwargs, None, None, None, None)
+
+        self.nonlin3 = nonlin(**nonlin_kwargs) if nonlin is not None else lambda x: x
+
+        # Stochastic Depth
+        self.apply_stochastic_depth = False if stochastic_depth_p == 0.0 else True
+        if self.apply_stochastic_depth:
+            self.drop_path = DropPath(drop_prob=stochastic_depth_p)
+
+        # Squeeze Excitation
+        self.apply_se = squeeze_excitation
+        if self.apply_se:
+            self.squeeze_excitation = SqueezeExcite(self.output_channels, conv_op,
+                                                    rd_ratio=squeeze_excitation_reduction_ratio, rd_divisor=8)
+
+        has_stride = (isinstance(stride, int) and stride != 1) or any([i != 1 for i in stride])
+        requires_projection = (input_channels != output_channels)
+
+        if has_stride or requires_projection:
+            ops = []
+            if has_stride:
+                ops.append(get_matching_pool_op(conv_op=conv_op, adaptive=False, pool_type='avg')(stride, stride))
+            if requires_projection:
+                ops.append(
+                    ConvDropoutNormReLU(conv_op, input_channels, output_channels, 1, 1, False,
+                                        norm_op, norm_op_kwargs, None, None, None, None
+                                        )
+                )
+            self.skip = nn.Sequential(*ops)
+        else:
+            self.skip = lambda x: x
+
+    def forward(self, x):
+        residual = self.skip(x)
+        out = self.conv3(self.conv2(self.conv1(x)))
+        if self.apply_stochastic_depth:
+            out = self.drop_path(out)
+        if self.apply_se:
+            out = self.squeeze_excitation(out)
+        out += residual
+        return self.nonlin3(out)
+
+    def compute_conv_feature_map_size(self, input_size):
+        assert len(input_size) == len(self.stride), "just give the image size without color/feature channels or " \
+                                                    "batch channel. Do not give input_size=(b, c, x, y(, z)). " \
+                                                    "Give input_size=(x, y(, z))!"
+        size_after_stride = [i // j for i, j in zip(input_size, self.stride)]
+        # conv1
+        output_size_conv1 = np.prod([self.bottleneck_channels, *input_size], dtype=np.int64)
+        # conv2
+        output_size_conv2 = np.prod([self.bottleneck_channels, *size_after_stride], dtype=np.int64)
+        # conv3
+        output_size_conv3 = np.prod([self.output_channels, *size_after_stride], dtype=np.int64)
+        # skip conv (if applicable)
+        if (self.input_channels != self.output_channels) or any([i != j for i, j in zip(input_size, size_after_stride)]):
+            assert isinstance(self.skip, nn.Sequential)
+            output_size_skip = np.prod([self.output_channels, *size_after_stride], dtype=np.int64)
+        else:
+            assert not isinstance(self.skip, nn.Sequential)
+            output_size_skip = 0
+        return output_size_conv1 + output_size_conv2 + output_size_conv3 + output_size_skip
+
+
+class StackedResidualBlocks(nn.Module):
+    def __init__(self,
+                 n_blocks: int,
+                 conv_op: Type[_ConvNd],
+                 input_channels: int,
+                 output_channels: Union[int, List[int], Tuple[int, ...]],
+                 kernel_size: Union[int, List[int], Tuple[int, ...]],
+                 initial_stride: Union[int, List[int], Tuple[int, ...]],
+                 conv_bias: bool = False,
+                 norm_op: Union[None, Type[nn.Module]] = None,
+                 norm_op_kwargs: dict = None,
+                 dropout_op: Union[None, Type[_DropoutNd]] = None,
+                 dropout_op_kwargs: dict = None,
+                 nonlin: Union[None, Type[torch.nn.Module]] = None,
+                 nonlin_kwargs: dict = None,
+                 block: Union[Type[BasicBlockD], Type[BottleneckD]] = BasicBlockD,
+                 bottleneck_channels: Union[int, List[int], Tuple[int, ...]] = None,
+                 stochastic_depth_p: float = 0.0,
+                 squeeze_excitation: bool = False,
+                 squeeze_excitation_reduction_ratio: float = 1. / 16
+                 ):
+        """
+        Stack multiple instances of block.
+
+        :param n_blocks: number of residual blocks
+        :param conv_op: nn.ConvNd class
+        :param input_channels: only relevant for forst block in the sequence. This is the input number of features.
+        After the first block, the number of features in the main path to which the residuals are added is output_channels
+        :param output_channels: number of features in the main path to which the residuals are added (and also the
+        number of features of the output)
+        :param kernel_size: kernel size for all nxn (n!=1) convolutions. Default: 3x3
+        :param initial_stride: only affects the first block. All subsequent blocks have stride 1
+        :param conv_bias: usually False
+        :param norm_op: nn.BatchNormNd, InstanceNormNd etc
+        :param norm_op_kwargs: dictionary of kwargs. Leave empty ({}) for defaults
+        :param dropout_op: nn.DropoutNd, can be None for no dropout
+        :param dropout_op_kwargs:
+        :param nonlin:
+        :param nonlin_kwargs:
+        :param block: BasicBlockD or BottleneckD
+        :param bottleneck_channels: if block is BottleneckD then we need to know the number of bottleneck features.
+        Bottleneck will use first 1x1 conv to reduce input to bottleneck features, then run the nxn (see kernel_size)
+        conv on that (bottleneck -> bottleneck). Finally the output will be projected back to output_channels
+        (bottleneck -> output_channels) with the final 1x1 conv
+        :param stochastic_depth_p: probability of applying stochastic depth in residual blocks
+        :param squeeze_excitation: whether to apply squeeze and excitation or not
+        :param squeeze_excitation_reduction_ratio: ratio by how much squeeze and excitation should reduce channels
+        respective to number of out channels of respective block
+        """
+        super().__init__()
+        assert n_blocks > 0, 'n_blocks must be > 0'
+        assert block in [BasicBlockD, BottleneckD], 'block must be BasicBlockD or BottleneckD'
+        if not isinstance(output_channels, (tuple, list)):
+            output_channels = [output_channels] * n_blocks
+        if not isinstance(bottleneck_channels, (tuple, list)):
+            bottleneck_channels = [bottleneck_channels] * n_blocks
+
+        if block == BasicBlockD:
+            blocks = nn.Sequential(
+                block(conv_op, input_channels, output_channels[0], kernel_size, initial_stride, conv_bias,
+                      norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs, nonlin, nonlin_kwargs, stochastic_depth_p,
+                      squeeze_excitation, squeeze_excitation_reduction_ratio),
+                *[block(conv_op, output_channels[n - 1], output_channels[n], kernel_size, 1, conv_bias, norm_op,
+                        norm_op_kwargs, dropout_op, dropout_op_kwargs, nonlin, nonlin_kwargs, stochastic_depth_p,
+                        squeeze_excitation, squeeze_excitation_reduction_ratio) for n in range(1, n_blocks)]
+            )
+        else:
+            blocks = nn.Sequential(
+                block(conv_op, input_channels, bottleneck_channels[0], output_channels[0], kernel_size,
+                      initial_stride, conv_bias, norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs,
+                      nonlin, nonlin_kwargs, stochastic_depth_p, squeeze_excitation, squeeze_excitation_reduction_ratio),
+                *[block(conv_op, output_channels[n - 1], bottleneck_channels[n], output_channels[n], kernel_size,
+                        1, conv_bias, norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs,
+                        nonlin, nonlin_kwargs, stochastic_depth_p, squeeze_excitation,
+                        squeeze_excitation_reduction_ratio) for n in range(1, n_blocks)]
+            )
+        self.blocks = blocks
+        self.initial_stride = maybe_convert_scalar_to_list(conv_op, initial_stride)
+        self.output_channels = output_channels[-1]
+
+    def forward(self, x):
+        return self.blocks(x)
+
+    def compute_conv_feature_map_size(self, input_size):
+        assert len(input_size) == len(self.initial_stride), "just give the image size without color/feature channels or " \
+                                                    "batch channel. Do not give input_size=(b, c, x, y(, z)). " \
+                                                    "Give input_size=(x, y(, z))!"
+        output = self.blocks[0].compute_conv_feature_map_size(input_size)
+        size_after_stride = [i // j for i, j in zip(input_size, self.initial_stride)]
+        for b in self.blocks[1:]:
+            output += b.compute_conv_feature_map_size(size_after_stride)
+        return output
+
+
+if __name__ == '__main__':
+    data = torch.rand((1, 3, 40, 32))
+
+    stx = StackedResidualBlocks(2, nn.Conv2d, 24, (16, 16), (3, 3), (1, 2),
+                                                norm_op=nn.BatchNorm2d, nonlin=nn.ReLU, nonlin_kwargs={'inplace': True},
+                                                block=BottleneckD, bottleneck_channels=3)
+    model = nn.Sequential(ConvDropoutNormReLU(nn.Conv2d,
+                                              3, 24, 3, 1, True, nn.BatchNorm2d, {}, None, None, nn.LeakyReLU,
+                                              {'inplace': True}),
+                          stx)
+    import hiddenlayer as hl
+
+    g = hl.build_graph(model, data,
+                       transforms=None)
+    g.save("network_architecture.pdf")
+    del g
+
+    print(stx.compute_conv_feature_map_size((40, 32)))

RADAR_inference/dynamic_network_architectures/building_blocks/residual_encoders.py

@@ -0,0 +1,172 @@
+import torch
+from torch import nn
+import numpy as np
+from typing import Union, Type, List, Tuple
+
+from torch.nn.modules.conv import _ConvNd
+from torch.nn.modules.dropout import _DropoutNd
+from dynamic_network_architectures.building_blocks.residual import StackedResidualBlocks, BottleneckD, BasicBlockD
+from dynamic_network_architectures.building_blocks.helper import maybe_convert_scalar_to_list, get_matching_pool_op
+from dynamic_network_architectures.building_blocks.simple_conv_blocks import StackedConvBlocks
+
+
+class ResidualEncoder(nn.Module):
+    def __init__(self,
+                 input_channels: int,
+                 n_stages: int,
+                 features_per_stage: Union[int, List[int], Tuple[int, ...]],
+                 conv_op: Type[_ConvNd],
+                 kernel_sizes: Union[int, List[int], Tuple[int, ...]],
+                 strides: Union[int, List[int], Tuple[int, ...], Tuple[Tuple[int, ...], ...]],
+                 n_blocks_per_stage: Union[int, List[int], Tuple[int, ...]],
+                 conv_bias: bool = False,
+                 norm_op: Union[None, Type[nn.Module]] = None,
+                 norm_op_kwargs: dict = None,
+                 dropout_op: Union[None, Type[_DropoutNd]] = None,
+                 dropout_op_kwargs: dict = None,
+                 nonlin: Union[None, Type[torch.nn.Module]] = None,
+                 nonlin_kwargs: dict = None,
+                 block: Union[Type[BasicBlockD], Type[BottleneckD]] = BasicBlockD,
+                 bottleneck_channels: Union[int, List[int], Tuple[int, ...]] = None,
+                 return_skips: bool = False,
+                 disable_default_stem: bool = False,
+                 stem_channels: int = None,
+                 pool_type: str = 'conv',
+                 stochastic_depth_p: float = 0.0,
+                 squeeze_excitation: bool = False,
+                 squeeze_excitation_reduction_ratio: float = 1. / 16
+                 ):
+        """
+
+        :param input_channels:
+        :param n_stages:
+        :param features_per_stage: Note: If the block is BottleneckD, then this number is supposed to be the number of
+        features AFTER the expansion (which is not coded implicitly in this repository)! See todo!
+        :param conv_op:
+        :param kernel_sizes:
+        :param strides:
+        :param n_blocks_per_stage:
+        :param conv_bias:
+        :param norm_op:
+        :param norm_op_kwargs:
+        :param dropout_op:
+        :param dropout_op_kwargs:
+        :param nonlin:
+        :param nonlin_kwargs:
+        :param block:
+        :param bottleneck_channels: only needed if block is BottleneckD
+        :param return_skips: set this to True if used as encoder in a U-Net like network
+        :param disable_default_stem: If True then no stem will be created. You need to build your own and ensure it is executed first, see todo.
+        The stem in this implementation does not so stride/pooling so building your own stem is a necessity if you need this.
+        :param stem_channels: if None, features_per_stage[0] will be used for the default stem. Not recommended for BottleneckD
+        :param pool_type: if conv, strided conv will be used. avg = average pooling, max = max pooling
+        """
+        super().__init__()
+        if isinstance(kernel_sizes, int):
+            kernel_sizes = [kernel_sizes] * n_stages
+        if isinstance(features_per_stage, int):
+            features_per_stage = [features_per_stage] * n_stages
+        if isinstance(n_blocks_per_stage, int):
+            n_blocks_per_stage = [n_blocks_per_stage] * n_stages
+        if isinstance(strides, int):
+            strides = [strides] * n_stages
+        if bottleneck_channels is None or isinstance(bottleneck_channels, int):
+            bottleneck_channels = [bottleneck_channels] * n_stages
+        assert len(
+            bottleneck_channels) == n_stages, "bottleneck_channels must be None or have as many entries as we have resolution stages (n_stages)"
+        assert len(
+            kernel_sizes) == n_stages, "kernel_sizes must have as many entries as we have resolution stages (n_stages)"
+        assert len(
+            n_blocks_per_stage) == n_stages, "n_conv_per_stage must have as many entries as we have resolution stages (n_stages)"
+        assert len(
+            features_per_stage) == n_stages, "features_per_stage must have as many entries as we have resolution stages (n_stages)"
+        assert len(strides) == n_stages, "strides must have as many entries as we have resolution stages (n_stages). " \
+                                         "Important: first entry is recommended to be 1, else we run strided conv drectly on the input"
+
+        pool_op = get_matching_pool_op(conv_op, pool_type=pool_type) if pool_type != 'conv' else None
+
+        # build a stem, Todo maybe we need more flexibility for this in the future. For now, if you need a custom
+        #  stem you can just disable the stem and build your own.
+        #  THE STEM DOES NOT DO STRIDE/POOLING IN THIS IMPLEMENTATION
+        if not disable_default_stem:
+            if stem_channels is None:
+                stem_channels = features_per_stage[0]
+            self.stem = StackedConvBlocks(1, conv_op, input_channels, stem_channels, kernel_sizes[0], 1, conv_bias,
+                                          norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs, nonlin, nonlin_kwargs)
+            input_channels = stem_channels
+        else:
+            self.stem = None
+
+        # now build the network
+        stages = []
+        for s in range(n_stages):
+            stride_for_conv = strides[s] if pool_op is None else 1
+
+            stage = StackedResidualBlocks(
+                n_blocks_per_stage[s], conv_op, input_channels, features_per_stage[s], kernel_sizes[s], stride_for_conv,
+                conv_bias, norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs, nonlin, nonlin_kwargs,
+                block=block, bottleneck_channels=bottleneck_channels[s], stochastic_depth_p=stochastic_depth_p,
+                squeeze_excitation=squeeze_excitation,
+                squeeze_excitation_reduction_ratio=squeeze_excitation_reduction_ratio
+            )
+
+            if pool_op is not None:
+                stage = nn.Sequential(pool_op(strides[s]), stage)
+
+            stages.append(stage)
+            input_channels = features_per_stage[s]
+
+        self.stages = nn.Sequential(*stages)
+        self.output_channels = features_per_stage
+        self.strides = [maybe_convert_scalar_to_list(conv_op, i) for i in strides]
+        self.return_skips = return_skips
+
+        # we store some things that a potential decoder needs
+        self.conv_op = conv_op
+        self.norm_op = norm_op
+        self.norm_op_kwargs = norm_op_kwargs
+        self.nonlin = nonlin
+        self.nonlin_kwargs = nonlin_kwargs
+        self.dropout_op = dropout_op
+        self.dropout_op_kwargs = dropout_op_kwargs
+        self.conv_bias = conv_bias
+        self.kernel_sizes = kernel_sizes
+
+    def forward(self, x):
+        if self.stem is not None:
+            x = self.stem(x)
+        ret = []
+        for s in self.stages:
+            x = s(x)
+            ret.append(x)
+        if self.return_skips:
+            return ret
+        else:
+            return ret[-1]
+
+    def compute_conv_feature_map_size(self, input_size):
+        if self.stem is not None:
+            output = self.stem.compute_conv_feature_map_size(input_size)
+        else:
+            output = np.int64(0)
+
+        for s in range(len(self.stages)):
+            output += self.stages[s].compute_conv_feature_map_size(input_size)
+            input_size = [i // j for i, j in zip(input_size, self.strides[s])]
+
+        return output
+
+
+if __name__ == '__main__':
+    data = torch.rand((1, 3, 128, 160))
+
+    model = ResidualEncoder(3, 5, (2, 4, 6, 8, 10), nn.Conv2d, 3, ((1, 1), 2, (2, 2), (2, 2), (2, 2)), 2, False,
+                            nn.BatchNorm2d, None, None, None, nn.ReLU, None, stem_channels=7)
+    import hiddenlayer as hl
+
+    g = hl.build_graph(model, data,
+                       transforms=None)
+    g.save("network_architecture.pdf")
+    del g
+
+    print(model.compute_conv_feature_map_size((128, 160)))

RADAR_inference/dynamic_network_architectures/building_blocks/simple_conv_blocks.py

@@ -0,0 +1,167 @@
+from typing import Tuple, List, Union, Type
+
+import numpy as np
+import torch.nn
+from torch import nn
+from torch.nn.modules.conv import _ConvNd
+from torch.nn.modules.dropout import _DropoutNd
+
+from dynamic_network_architectures.building_blocks.helper import maybe_convert_scalar_to_list
+
+
+class ConvDropoutNormReLU(nn.Module):
+    def __init__(self,
+                 conv_op: Type[_ConvNd],
+                 input_channels: int,
+                 output_channels: int,
+                 kernel_size: Union[int, List[int], Tuple[int, ...]],
+                 stride: Union[int, List[int], Tuple[int, ...]],
+                 conv_bias: bool = False,
+                 norm_op: Union[None, Type[nn.Module]] = None,
+                 norm_op_kwargs: dict = None,
+                 dropout_op: Union[None, Type[_DropoutNd]] = None,
+                 dropout_op_kwargs: dict = None,
+                 nonlin: Union[None, Type[torch.nn.Module]] = None,
+                 nonlin_kwargs: dict = None,
+                 nonlin_first: bool = False
+                 ):
+        super(ConvDropoutNormReLU, self).__init__()
+        self.input_channels = input_channels
+        self.output_channels = output_channels
+        stride = maybe_convert_scalar_to_list(conv_op, stride)
+        self.stride = stride
+
+        kernel_size = maybe_convert_scalar_to_list(conv_op, kernel_size)
+        if norm_op_kwargs is None:
+            norm_op_kwargs = {}
+        if nonlin_kwargs is None:
+            nonlin_kwargs = {}
+
+        ops = []
+
+        self.conv = conv_op(
+            input_channels,
+            output_channels,
+            kernel_size,
+            stride,
+            padding=[(i - 1) // 2 for i in kernel_size],
+            dilation=1,
+            bias=conv_bias,
+        )
+        ops.append(self.conv)
+
+        if dropout_op is not None:
+            self.dropout = dropout_op(**dropout_op_kwargs)
+            ops.append(self.dropout)
+
+        if norm_op is not None:
+            self.norm = norm_op(output_channels, **norm_op_kwargs)
+            ops.append(self.norm)
+
+        if nonlin is not None:
+            self.nonlin = nonlin(**nonlin_kwargs)
+            ops.append(self.nonlin)
+
+        if nonlin_first and (norm_op is not None and nonlin is not None):
+            ops[-1], ops[-2] = ops[-2], ops[-1]
+
+        self.all_modules = nn.Sequential(*ops)
+
+    def forward(self, x):
+        return self.all_modules(x)
+
+    def compute_conv_feature_map_size(self, input_size):
+        assert len(input_size) == len(self.stride), "just give the image size without color/feature channels or " \
+                                                    "batch channel. Do not give input_size=(b, c, x, y(, z)). " \
+                                                    "Give input_size=(x, y(, z))!"
+        output_size = [i // j for i, j in zip(input_size, self.stride)]  # we always do same padding
+        return np.prod([self.output_channels, *output_size], dtype=np.int64)
+
+
+class StackedConvBlocks(nn.Module):
+    def __init__(self,
+                 num_convs: int,
+                 conv_op: Type[_ConvNd],
+                 input_channels: int,
+                 output_channels: Union[int, List[int], Tuple[int, ...]],
+                 kernel_size: Union[int, List[int], Tuple[int, ...]],
+                 initial_stride: Union[int, List[int], Tuple[int, ...]],
+                 conv_bias: bool = False,
+                 norm_op: Union[None, Type[nn.Module]] = None,
+                 norm_op_kwargs: dict = None,
+                 dropout_op: Union[None, Type[_DropoutNd]] = None,
+                 dropout_op_kwargs: dict = None,
+                 nonlin: Union[None, Type[torch.nn.Module]] = None,
+                 nonlin_kwargs: dict = None,
+                 nonlin_first: bool = False
+                 ):
+        """
+
+        :param conv_op:
+        :param num_convs:
+        :param input_channels:
+        :param output_channels: can be int or a list/tuple of int. If list/tuple are provided, each entry is for
+        one conv. The length of the list/tuple must then naturally be num_convs
+        :param kernel_size:
+        :param initial_stride:
+        :param conv_bias:
+        :param norm_op:
+        :param norm_op_kwargs:
+        :param dropout_op:
+        :param dropout_op_kwargs:
+        :param nonlin:
+        :param nonlin_kwargs:
+        """
+        super().__init__()
+        if not isinstance(output_channels, (tuple, list)):
+            output_channels = [output_channels] * num_convs
+
+        self.convs = nn.Sequential(
+            ConvDropoutNormReLU(
+                conv_op, input_channels, output_channels[0], kernel_size, initial_stride, conv_bias, norm_op,
+                norm_op_kwargs, dropout_op, dropout_op_kwargs, nonlin, nonlin_kwargs, nonlin_first
+            ),
+            *[
+                ConvDropoutNormReLU(
+                    conv_op, output_channels[i - 1], output_channels[i], kernel_size, 1, conv_bias, norm_op,
+                    norm_op_kwargs, dropout_op, dropout_op_kwargs, nonlin, nonlin_kwargs, nonlin_first
+                )
+                for i in range(1, num_convs)
+            ]
+        )
+
+        self.output_channels = output_channels[-1]
+        self.initial_stride = maybe_convert_scalar_to_list(conv_op, initial_stride)
+
+    def forward(self, x):
+        return self.convs(x)
+
+    def compute_conv_feature_map_size(self, input_size):
+        assert len(input_size) == len(self.initial_stride), "just give the image size without color/feature channels or " \
+                                                            "batch channel. Do not give input_size=(b, c, x, y(, z)). " \
+                                                            "Give input_size=(x, y(, z))!"
+        output = self.convs[0].compute_conv_feature_map_size(input_size)
+        size_after_stride = [i // j for i, j in zip(input_size, self.initial_stride)]
+        for b in self.convs[1:]:
+            output += b.compute_conv_feature_map_size(size_after_stride)
+        return output
+
+
+if __name__ == '__main__':
+    data = torch.rand((1, 3, 40, 32))
+
+    stx = StackedConvBlocks(2, nn.Conv2d, 24, 16, (3, 3), 2,
+                            norm_op=nn.BatchNorm2d, nonlin=nn.ReLU, nonlin_kwargs={'inplace': True},
+                            )
+    model = nn.Sequential(ConvDropoutNormReLU(nn.Conv2d,
+                                              3, 24, 3, 1, True, nn.BatchNorm2d, {}, None, None, nn.LeakyReLU,
+                                              {'inplace': True}),
+                          stx)
+    import hiddenlayer as hl
+
+    g = hl.build_graph(model, data,
+                       transforms=None)
+    g.save("network_architecture.pdf")
+    del g
+
+    stx.compute_conv_feature_map_size((40, 32))

RADAR_inference/dynamic_network_architectures/building_blocks/unet_decoder.py

@@ -0,0 +1,154 @@
+import numpy as np
+import torch
+from torch import nn
+from typing import Union, List, Tuple, Type
+
+from torch.nn.modules.dropout import _DropoutNd
+
+from dynamic_network_architectures.building_blocks.simple_conv_blocks import StackedConvBlocks
+from dynamic_network_architectures.building_blocks.helper import get_matching_convtransp
+from dynamic_network_architectures.building_blocks.residual_encoders import ResidualEncoder
+from dynamic_network_architectures.building_blocks.plain_conv_encoder import PlainConvEncoder
+
+
+class UNetDecoder(nn.Module):
+    def __init__(self,
+                 encoder: Union[PlainConvEncoder, ResidualEncoder],
+                 num_classes: int,
+                 n_conv_per_stage: Union[int, Tuple[int, ...], List[int]],
+                 deep_supervision,
+                 nonlin_first: bool = False,
+                 norm_op: Union[None, Type[nn.Module]] = None,
+                 norm_op_kwargs: dict = None,
+                 dropout_op: Union[None, Type[_DropoutNd]] = None,
+                 dropout_op_kwargs: dict = None,
+                 nonlin: Union[None, Type[torch.nn.Module]] = None,
+                 nonlin_kwargs: dict = None,
+                 conv_bias: bool = None
+                 ):
+        """
+        This class needs the skips of the encoder as input in its forward.
+
+        the encoder goes all the way to the bottleneck, so that's where the decoder picks up. stages in the decoder
+        are sorted by order of computation, so the first stage has the lowest resolution and takes the bottleneck
+        features and the lowest skip as inputs
+        the decoder has two (three) parts in each stage:
+        1) conv transpose to upsample the feature maps of the stage below it (or the bottleneck in case of the first stage)
+        2) n_conv_per_stage conv blocks to let the two inputs get to know each other and merge
+        3) (optional if deep_supervision=True) a segmentation output Todo: enable upsample logits?
+        :param encoder:
+        :param num_classes:
+        :param n_conv_per_stage:
+        :param deep_supervision:
+        """
+        super().__init__()
+        self.deep_supervision = deep_supervision
+        self.encoder = encoder
+        self.num_classes = num_classes
+        n_stages_encoder = len(encoder.output_channels)
+        if isinstance(n_conv_per_stage, int):
+            n_conv_per_stage = [n_conv_per_stage] * (n_stages_encoder - 1)
+        assert len(n_conv_per_stage) == n_stages_encoder - 1, "n_conv_per_stage must have as many entries as we have " \
+                                                          "resolution stages - 1 (n_stages in encoder - 1), " \
+                                                          "here: %d" % n_stages_encoder
+
+        transpconv_op = get_matching_convtransp(conv_op=encoder.conv_op)
+        conv_bias = encoder.conv_bias if conv_bias is None else conv_bias
+        norm_op = encoder.norm_op if norm_op is None else norm_op
+        norm_op_kwargs = encoder.norm_op_kwargs if norm_op_kwargs is None else norm_op_kwargs
+        dropout_op = encoder.dropout_op if dropout_op is None else dropout_op
+        dropout_op_kwargs = encoder.dropout_op_kwargs if dropout_op_kwargs is None else dropout_op_kwargs
+        nonlin = encoder.nonlin if nonlin is None else nonlin
+        nonlin_kwargs = encoder.nonlin_kwargs if nonlin_kwargs is None else nonlin_kwargs
+
+
+        # we start with the bottleneck and work out way up
+        stages = []
+        transpconvs = []
+        seg_layers = []
+        for s in range(1, n_stages_encoder):
+            input_features_below = encoder.output_channels[-s]
+            input_features_skip = encoder.output_channels[-(s + 1)]
+            stride_for_transpconv = encoder.strides[-s]
+            transpconvs.append(transpconv_op(
+                input_features_below, input_features_skip, stride_for_transpconv, stride_for_transpconv,
+                bias=conv_bias
+            ))
+            # input features to conv is 2x input_features_skip (concat input_features_skip with transpconv output)
+            stages.append(StackedConvBlocks(
+                n_conv_per_stage[s-1], encoder.conv_op, 2 * input_features_skip, input_features_skip,
+                encoder.kernel_sizes[-(s + 1)], 1,
+                conv_bias,
+                norm_op,
+                norm_op_kwargs,
+                dropout_op,
+                dropout_op_kwargs,
+                nonlin,
+                nonlin_kwargs,
+                nonlin_first
+            ))
+
+            # we always build the deep supervision outputs so that we can always load parameters. If we don't do this
+            # then a model trained with deep_supervision=True could not easily be loaded at inference time where
+            # deep supervision is not needed. It's just a convenience thing
+            seg_layers.append(encoder.conv_op(input_features_skip, num_classes, 1, 1, 0, bias=True))
+
+        self.stages = nn.ModuleList(stages)
+        self.transpconvs = nn.ModuleList(transpconvs)
+        self.seg_layers = nn.ModuleList(seg_layers)
+
+    def forward(self, skips):
+        """
+        we expect to get the skips in the order they were computed, so the bottleneck should be the last entry
+        :param skips:
+        :return:
+        """
+        lres_input = skips[-1]
+        seg_outputs = []
+        for s in range(len(self.stages)):
+            x = self.transpconvs[s](lres_input)
+            x = torch.cat((x, skips[-(s+2)]), 1)
+            x = self.stages[s](x)
+            if self.deep_supervision:
+                seg_outputs.append(self.seg_layers[s](x))
+            elif s == (len(self.stages) - 1):
+                seg_outputs.append(self.seg_layers[-1](x))
+            lres_input = x
+
+        # invert seg outputs so that the largest segmentation prediction is returned first
+        seg_outputs = seg_outputs[::-1]
+
+        if not self.deep_supervision:
+            r = seg_outputs[0]
+        else:
+            r = seg_outputs
+        return r
+
+    def compute_conv_feature_map_size(self, input_size):
+        """
+        IMPORTANT: input_size is the input_size of the encoder!
+        :param input_size:
+        :return:
+        """
+        # first we need to compute the skip sizes. Skip bottleneck because all output feature maps of our ops will at
+        # least have the size of the skip above that (therefore -1)
+        skip_sizes = []
+        for s in range(len(self.encoder.strides) - 1):
+            skip_sizes.append([i // j for i, j in zip(input_size, self.encoder.strides[s])])
+            input_size = skip_sizes[-1]
+        # print(skip_sizes)
+
+        assert len(skip_sizes) == len(self.stages)
+
+        # our ops are the other way around, so let's match things up
+        output = np.int64(0)
+        for s in range(len(self.stages)):
+            # print(skip_sizes[-(s+1)], self.encoder.output_channels[-(s+2)])
+            # conv blocks
+            output += self.stages[s].compute_conv_feature_map_size(skip_sizes[-(s+1)])
+            # trans conv
+            output += np.prod([self.encoder.output_channels[-(s+2)], *skip_sizes[-(s+1)]], dtype=np.int64)
+            # segmentation
+            if self.deep_supervision or (s == (len(self.stages) - 1)):
+                output += np.prod([self.num_classes, *skip_sizes[-(s+1)]], dtype=np.int64)
+        return output

RADAR_inference/dynamic_network_architectures/building_blocks/unet_decoder_light.py

@@ -0,0 +1,154 @@
+import numpy as np
+import torch
+from torch import nn
+from typing import Union, List, Tuple, Type
+
+from torch.nn.modules.dropout import _DropoutNd
+
+from dynamic_network_architectures.building_blocks.simple_conv_blocks import StackedConvBlocks
+from dynamic_network_architectures.building_blocks.helper import get_matching_convtransp
+from dynamic_network_architectures.building_blocks.residual_encoders import ResidualEncoder
+from dynamic_network_architectures.building_blocks.plain_conv_encoder import PlainConvEncoder
+
+
+class UNetDecoder(nn.Module):
+    def __init__(self,
+                 encoder: Union[PlainConvEncoder, ResidualEncoder],
+                 num_classes: int,
+                 n_conv_per_stage: Union[int, Tuple[int, ...], List[int]],
+                 deep_supervision,
+                 nonlin_first: bool = False,
+                 norm_op: Union[None, Type[nn.Module]] = None,
+                 norm_op_kwargs: dict = None,
+                 dropout_op: Union[None, Type[_DropoutNd]] = None,
+                 dropout_op_kwargs: dict = None,
+                 nonlin: Union[None, Type[torch.nn.Module]] = None,
+                 nonlin_kwargs: dict = None,
+                 conv_bias: bool = None
+                 ):
+        """
+        This class needs the skips of the encoder as input in its forward.
+
+        the encoder goes all the way to the bottleneck, so that's where the decoder picks up. stages in the decoder
+        are sorted by order of computation, so the first stage has the lowest resolution and takes the bottleneck
+        features and the lowest skip as inputs
+        the decoder has two (three) parts in each stage:
+        1) conv transpose to upsample the feature maps of the stage below it (or the bottleneck in case of the first stage)
+        2) n_conv_per_stage conv blocks to let the two inputs get to know each other and merge
+        3) (optional if deep_supervision=True) a segmentation output Todo: enable upsample logits?
+        :param encoder:
+        :param num_classes:
+        :param n_conv_per_stage:
+        :param deep_supervision:
+        """
+        super().__init__()
+        self.deep_supervision = deep_supervision
+        self.encoder = encoder
+        self.num_classes = num_classes
+        n_stages_encoder = len(encoder.output_channels)
+        if isinstance(n_conv_per_stage, int):
+            n_conv_per_stage = [n_conv_per_stage] * (n_stages_encoder - 1)
+        # assert len(n_conv_per_stage) == n_stages_encoder - 1, "n_conv_per_stage must have as many entries as we have " \
+        #                                                   "resolution stages - 1 (n_stages in encoder - 1), " \
+        #                                                   "here: %d" % n_stages_encoder
+
+        transpconv_op = get_matching_convtransp(conv_op=encoder.conv_op)
+        conv_bias = encoder.conv_bias if conv_bias is None else conv_bias
+        norm_op = encoder.norm_op if norm_op is None else norm_op
+        norm_op_kwargs = encoder.norm_op_kwargs if norm_op_kwargs is None else norm_op_kwargs
+        dropout_op = encoder.dropout_op if dropout_op is None else dropout_op
+        dropout_op_kwargs = encoder.dropout_op_kwargs if dropout_op_kwargs is None else dropout_op_kwargs
+        nonlin = encoder.nonlin if nonlin is None else nonlin
+        nonlin_kwargs = encoder.nonlin_kwargs if nonlin_kwargs is None else nonlin_kwargs
+
+
+        stages = []
+        transpconvs = []
+        seg_layers = []
+        for s in range(1, n_stages_encoder-1):
+            input_features_below = encoder.output_channels[-s]
+            input_features_skip = encoder.output_channels[-(s + 1)]
+            stride_for_transpconv = encoder.strides[-s]
+            transpconvs.append(transpconv_op(
+                input_features_below, input_features_skip, stride_for_transpconv, stride_for_transpconv,
+                bias=conv_bias
+            ))
+            # input features to conv is 2x input_features_skip (concat input_features_skip with transpconv output)
+            stages.append(StackedConvBlocks(
+                n_conv_per_stage[s-1], encoder.conv_op, 2 * input_features_skip, input_features_skip,
+                encoder.kernel_sizes[-(s + 1)], 1,
+                conv_bias,
+                norm_op,
+                norm_op_kwargs,
+                dropout_op,
+                dropout_op_kwargs,
+                nonlin,
+                nonlin_kwargs,
+                nonlin_first
+            ))
+
+            # we always build the deep supervision outputs so that we can always load parameters. If we don't do this
+            # then a model trained with deep_supervision=True could not easily be loaded at inference time where
+            # deep supervision is not needed. It's just a convenience thing
+            seg_layers.append(encoder.conv_op(input_features_skip, num_classes, 1, 1, 0, bias=True))
+
+        self.stages = nn.ModuleList(stages)
+        self.transpconvs = nn.ModuleList(transpconvs)
+        self.seg_layers = nn.ModuleList(seg_layers)
+
+    def forward(self, skips):
+        """
+        we expect to get the skips in the order they were computed, so the bottleneck should be the last entry
+        :param skips:
+        :return:
+        """
+        # print('debug light nnunvetv2 decoder')
+        lres_input = skips[-1]
+        seg_outputs = []
+        for s in range(len(self.stages)):
+            x = self.transpconvs[s](lres_input)
+            x = torch.cat((x, skips[-(s+2)]), 1)
+            x = self.stages[s](x)
+            if self.deep_supervision:
+                seg_outputs.append(self.seg_layers[s](x))
+            elif s == (len(self.stages) - 1):
+                seg_outputs.append(self.seg_layers[-1](x))
+            lres_input = x
+
+        # invert seg outputs so that the largest segmentation prediction is returned first
+        seg_outputs = seg_outputs[::-1]
+
+        if not self.deep_supervision:
+            r = seg_outputs[0]
+        else:
+            r = seg_outputs
+        return r
+
+    def compute_conv_feature_map_size(self, input_size):
+        """
+        IMPORTANT: input_size is the input_size of the encoder!
+        :param input_size:
+        :return:
+        """
+        # first we need to compute the skip sizes. Skip bottleneck because all output feature maps of our ops will at
+        # least have the size of the skip above that (therefore -1)
+        skip_sizes = []
+        for s in range(len(self.encoder.strides) - 1):
+            skip_sizes.append([i // j for i, j in zip(input_size, self.encoder.strides[s])])
+            input_size = skip_sizes[-1]
+        # print(skip_sizes)
+
+        assert len(skip_sizes) == len(self.stages)
+
+        # our ops are the other way around, so let's match things up
+        output = np.int64(0)
+        for s in range(len(self.stages)):
+            # print(skip_sizes[-(s+1)], self.encoder.output_channels[-(s+2)])
+            # conv blocks
+            output += self.stages[s].compute_conv_feature_map_size(skip_sizes[-(s+1)])
+            # trans conv
+            output += np.prod([self.encoder.output_channels[-(s+2)], *skip_sizes[-(s+1)]], dtype=np.int64)
+            # segmentation
+            if self.deep_supervision or (s == (len(self.stages) - 1)):
+                output += np.prod([self.num_classes, *skip_sizes[-(s+1)]], dtype=np.int64)
+        return output

RADAR_inference/dynamic_network_architectures/building_blocks/unet_residual_decoder.py

@@ -0,0 +1,155 @@
+import numpy as np
+from typing import Union, Tuple, List, Type
+
+import torch
+from dynamic_network_architectures.building_blocks.helper import get_matching_convtransp
+from dynamic_network_architectures.building_blocks.plain_conv_encoder import PlainConvEncoder
+from dynamic_network_architectures.building_blocks.residual import StackedResidualBlocks
+from dynamic_network_architectures.building_blocks.residual_encoders import ResidualEncoder
+from torch import nn
+from torch.nn.modules.dropout import _DropoutNd
+
+
+class UNetResDecoder(nn.Module):
+    def __init__(self,
+                 encoder: Union[PlainConvEncoder, ResidualEncoder],
+                 num_classes: int,
+                 n_conv_per_stage: Union[int, Tuple[int, ...], List[int]],
+                 deep_supervision,
+                 nonlin_first: bool = False,
+                 norm_op: Union[None, Type[nn.Module]] = None,
+                 norm_op_kwargs: dict = None,
+                 dropout_op: Union[None, Type[_DropoutNd]] = None,
+                 dropout_op_kwargs: dict = None,
+                 nonlin: Union[None, Type[torch.nn.Module]] = None,
+                 nonlin_kwargs: dict = None,
+                 conv_bias: bool = None
+                 ):
+        """
+        This class needs the skips of the encoder as input in its forward.
+
+        the encoder goes all the way to the bottleneck, so that's where the decoder picks up. stages in the decoder
+        are sorted by order of computation, so the first stage has the lowest resolution and takes the bottleneck
+        features and the lowest skip as inputs
+        the decoder has two (three) parts in each stage:
+        1) conv transpose to upsample the feature maps of the stage below it (or the bottleneck in case of the first stage)
+        2) n_conv_per_stage conv blocks to let the two inputs get to know each other and merge
+        3) (optional if deep_supervision=True) a segmentation output Todo: enable upsample logits?
+        :param encoder:
+        :param num_classes:
+        :param n_conv_per_stage:
+        :param deep_supervision:
+        """
+        super().__init__()
+        self.deep_supervision = deep_supervision
+        self.encoder = encoder
+        self.num_classes = num_classes
+        n_stages_encoder = len(encoder.output_channels)
+        if isinstance(n_conv_per_stage, int):
+            n_conv_per_stage = [n_conv_per_stage] * (n_stages_encoder - 1)
+        assert len(n_conv_per_stage) == n_stages_encoder - 1, "n_conv_per_stage must have as many entries as we have " \
+                                                              "resolution stages - 1 (n_stages in encoder - 1), " \
+                                                              "here: %d" % n_stages_encoder
+
+        transpconv_op = get_matching_convtransp(conv_op=encoder.conv_op)
+        conv_bias = encoder.conv_bias if conv_bias is None else conv_bias
+        norm_op = encoder.norm_op if norm_op is None else norm_op
+        norm_op_kwargs = encoder.norm_op_kwargs if norm_op_kwargs is None else norm_op_kwargs
+        dropout_op = encoder.dropout_op if dropout_op is None else dropout_op
+        dropout_op_kwargs = encoder.dropout_op_kwargs if dropout_op_kwargs is None else dropout_op_kwargs
+        nonlin = encoder.nonlin if nonlin is None else nonlin
+        nonlin_kwargs = encoder.nonlin_kwargs if nonlin_kwargs is None else nonlin_kwargs
+
+        # we start with the bottleneck and work out way up
+        stages = []
+        transpconvs = []
+        seg_layers = []
+        for s in range(1, n_stages_encoder):
+            input_features_below = encoder.output_channels[-s]
+            input_features_skip = encoder.output_channels[-(s + 1)]
+            stride_for_transpconv = encoder.strides[-s]
+            transpconvs.append(transpconv_op(
+                input_features_below, input_features_skip, stride_for_transpconv, stride_for_transpconv,
+                bias=encoder.conv_bias
+            ))
+            # input features to conv is 2x input_features_skip (concat input_features_skip with transpconv output)
+            stages.append(StackedResidualBlocks(
+                n_blocks=n_conv_per_stage[s - 1],
+                conv_op=encoder.conv_op,
+                input_channels=2 * input_features_skip,
+                output_channels=input_features_skip,
+                kernel_size=encoder.kernel_sizes[-(s + 1)],
+                initial_stride=1,
+                conv_bias=conv_bias,
+                norm_op=norm_op,
+                norm_op_kwargs=norm_op_kwargs,
+                dropout_op=dropout_op,
+                dropout_op_kwargs=dropout_op_kwargs,
+                nonlin=nonlin,
+                nonlin_kwargs=nonlin_kwargs,
+            ))
+
+            # we always build the deep supervision outputs so that we can always load parameters. If we don't do this
+            # then a model trained with deep_supervision=True could not easily be loaded at inference time where
+            # deep supervision is not needed. It's just a convenience thing
+            seg_layers.append(encoder.conv_op(input_features_skip, num_classes, 1, 1, 0, bias=True))
+
+        self.stages = nn.ModuleList(stages)
+        self.transpconvs = nn.ModuleList(transpconvs)
+        self.seg_layers = nn.ModuleList(seg_layers)
+
+    def forward(self, skips):
+        """
+        we expect to get the skips in the order they were computed, so the bottleneck should be the last entry
+        :param skips:
+        :return:
+        """
+        lres_input = skips[-1]
+        seg_outputs = []
+        for s in range(len(self.stages)):
+            x = self.transpconvs[s](lres_input)
+            x = torch.cat((x, skips[-(s + 2)]), 1)
+            x = self.stages[s](x)
+            if self.deep_supervision:
+                seg_outputs.append(self.seg_layers[s](x))
+            elif s == (len(self.stages) - 1):
+                seg_outputs.append(self.seg_layers[-1](x))
+            lres_input = x
+
+        # invert seg outputs so that the largest segmentation prediction is returned first
+        seg_outputs = seg_outputs[::-1]
+
+        if not self.deep_supervision:
+            r = seg_outputs[0]
+        else:
+            r = seg_outputs
+        return r
+
+    def compute_conv_feature_map_size(self, input_size):
+        """
+        IMPORTANT: input_size is the input_size of the encoder!
+        :param input_size:
+        :return:
+        """
+        # first we need to compute the skip sizes. Skip bottleneck because all output feature maps of our ops will at
+        # least have the size of the skip above that (therefore -1)
+        skip_sizes = []
+        for s in range(len(self.encoder.strides) - 1):
+            skip_sizes.append([i // j for i, j in zip(input_size, self.encoder.strides[s])])
+            input_size = skip_sizes[-1]
+        # print(skip_sizes)
+
+        assert len(skip_sizes) == len(self.stages)
+
+        # our ops are the other way around, so let's match things up
+        output = np.int64(0)
+        for s in range(len(self.stages)):
+            # print(skip_sizes[-(s+1)], self.encoder.output_channels[-(s+2)])
+            # conv blocks
+            output += self.stages[s].compute_conv_feature_map_size(skip_sizes[-(s + 1)])
+            # trans conv
+            output += np.prod([self.encoder.output_channels[-(s + 2)], *skip_sizes[-(s + 1)]], dtype=np.int64)
+            # segmentation
+            if self.deep_supervision or (s == (len(self.stages) - 1)):
+                output += np.prod([self.num_classes, *skip_sizes[-(s + 1)]], dtype=np.int64)
+        return output

RADAR_inference/dynamic_network_architectures/initialization/weight_init.py

@@ -0,0 +1,34 @@
+from torch import nn
+
+from dynamic_network_architectures.building_blocks.residual import BasicBlockD, BottleneckD
+
+
+class InitWeights_He(object):
+    def __init__(self, neg_slope: float = 1e-2):
+        self.neg_slope = neg_slope
+
+    def __call__(self, module):
+        if isinstance(module, nn.Conv3d) or isinstance(module, nn.Conv2d) or isinstance(module, nn.ConvTranspose2d) or isinstance(module, nn.ConvTranspose3d):
+            module.weight = nn.init.kaiming_normal_(module.weight, a=self.neg_slope)
+            if module.bias is not None:
+                module.bias = nn.init.constant_(module.bias, 0)
+
+
+class InitWeights_XavierUniform(object):
+    def __init__(self, gain: int = 1):
+        self.gain = gain
+
+    def __call__(self, module):
+        if isinstance(module, nn.Conv3d) or isinstance(module, nn.Conv2d) or isinstance(module, nn.ConvTranspose2d) or isinstance(module, nn.ConvTranspose3d):
+            module.weight = nn.init.xavier_uniform_(module.weight, self.gain)
+            if module.bias is not None:
+                module.bias = nn.init.constant_(module.bias, 0)
+
+
+def init_last_bn_before_add_to_0(module):
+    if isinstance(module, BasicBlockD):
+        module.conv2.norm.weight = nn.init.constant_(module.conv2.norm.weight, 0)
+        module.conv2.norm.bias = nn.init.constant_(module.conv2.norm.bias, 0)
+    if isinstance(module, BottleneckD):
+        module.conv3.norm.weight = nn.init.constant_(module.conv3.norm.weight, 0)
+        module.conv3.norm.bias = nn.init.constant_(module.conv3.norm.bias, 0)

RADAR_inference/dynamic_network_architectures/med.py

@@ -0,0 +1,1502 @@
+"""
+ Copyright (c) 2022, salesforce.com, inc.
+ All rights reserved.
+ SPDX-License-Identifier: BSD-3-Clause
+ For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
+ 
+ Based on huggingface code base
+ https://github.com/huggingface/transformers/blob/v4.15.0/src/transformers/models/bert
+"""
+
+import math
+import os
+import warnings
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor, device
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss
+import torch.nn.functional as F
+from transformers import BatchEncoding, PreTrainedTokenizer
+
+from transformers.activations import ACT2FN
+from transformers.file_utils import (
+    ModelOutput,
+)
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPastAndCrossAttentions,
+    BaseModelOutputWithPoolingAndCrossAttentions,
+    CausalLMOutputWithCrossAttentions,
+    MaskedLMOutput,
+    MultipleChoiceModelOutput,
+    NextSentencePredictorOutput,
+    QuestionAnsweringModelOutput,
+    SequenceClassifierOutput,
+    TokenClassifierOutput,
+)
+from transformers.modeling_utils import (
+    PreTrainedModel,
+    apply_chunking_to_forward,
+    find_pruneable_heads_and_indices,
+    prune_linear_layer,
+)
+from transformers.utils import logging
+from transformers.models.bert.configuration_bert import BertConfig
+
+# from lavis.models.base_model import BaseEncoder
+
+logging.set_verbosity_error()
+logger = logging.get_logger(__name__)
+
+
+class BaseEncoder(nn.Module):
+    """
+    Base class for primitive encoders, such as ViT, TimeSformer, etc.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward_features(self, samples, **kwargs):
+        raise NotImplementedError
+
+    @property
+    def device(self):
+        return list(self.parameters())[0].device
+    
+
+class BertEmbeddings(nn.Module):
+    """Construct the embeddings from word and position embeddings."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(
+            config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id
+        )
+        self.position_embeddings = nn.Embedding(
+            config.max_position_embeddings, config.hidden_size
+        )
+
+        if config.add_type_embeddings:
+            self.token_type_embeddings = nn.Embedding(
+                config.type_vocab_size, config.hidden_size
+            )
+
+        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
+        # any TensorFlow checkpoint file
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+        self.register_buffer(
+            "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1))
+        )
+        self.position_embedding_type = getattr(
+            config, "position_embedding_type", "absolute"
+        )
+
+        self.config = config
+
+    def forward(
+        self,
+        input_ids=None,
+        token_type_ids=None,
+        position_ids=None,
+        inputs_embeds=None,
+        past_key_values_length=0,
+    ):
+        if input_ids is not None:
+            input_shape = input_ids.size()
+        else:
+            input_shape = inputs_embeds.size()[:-1]
+
+        seq_length = input_shape[1]
+
+        if position_ids is None:
+            position_ids = self.position_ids[
+                :, past_key_values_length : seq_length + past_key_values_length
+            ]
+
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+
+        if token_type_ids is not None:
+            token_type_embeddings = self.token_type_embeddings(token_type_ids)
+
+            embeddings = inputs_embeds + token_type_embeddings
+        else:
+            embeddings = inputs_embeds
+
+        if self.position_embedding_type == "absolute":
+            position_embeddings = self.position_embeddings(position_ids)
+            embeddings += position_embeddings
+        embeddings = self.LayerNorm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+class BertSelfAttention(nn.Module):
+    def __init__(self, config, is_cross_attention):
+        super().__init__()
+        self.config = config
+        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
+            config, "embedding_size"
+        ):
+            raise ValueError(
+                "The hidden size (%d) is not a multiple of the number of attention "
+                "heads (%d)" % (config.hidden_size, config.num_attention_heads)
+            )
+
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        if is_cross_attention:
+            self.key = nn.Linear(config.encoder_width, self.all_head_size)
+            self.value = nn.Linear(config.encoder_width, self.all_head_size)
+        else:
+            self.key = nn.Linear(config.hidden_size, self.all_head_size)
+            self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+        self.position_embedding_type = getattr(
+            config, "position_embedding_type", "absolute"
+        )
+        if (
+            self.position_embedding_type == "relative_key"
+            or self.position_embedding_type == "relative_key_query"
+        ):
+            self.max_position_embeddings = config.max_position_embeddings
+            self.distance_embedding = nn.Embedding(
+                2 * config.max_position_embeddings - 1, self.attention_head_size
+            )
+        self.save_attention = False
+
+    def save_attn_gradients(self, attn_gradients):
+        self.attn_gradients = attn_gradients
+
+    def get_attn_gradients(self):
+        return self.attn_gradients
+
+    def save_attention_map(self, attention_map):
+        self.attention_map = attention_map
+
+    def get_attention_map(self):
+        return self.attention_map
+
+    def transpose_for_scores(self, x):
+        new_x_shape = x.size()[:-1] + (
+            self.num_attention_heads,
+            self.attention_head_size,
+        )
+        x = x.view(*new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+    ):
+        mixed_query_layer = self.query(hidden_states)
+
+        # If this is instantiated as a cross-attention module, the keys
+        # and values come from an encoder; the attention mask needs to be
+        # such that the encoder's padding tokens are not attended to.
+        is_cross_attention = encoder_hidden_states is not None
+
+        if is_cross_attention:
+            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
+            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
+            attention_mask = encoder_attention_mask
+        elif past_key_value is not None:
+            key_layer = self.transpose_for_scores(self.key(hidden_states))
+            value_layer = self.transpose_for_scores(self.value(hidden_states))
+            key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
+            value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
+        else:
+            key_layer = self.transpose_for_scores(self.key(hidden_states))
+            value_layer = self.transpose_for_scores(self.value(hidden_states))
+
+        query_layer = self.transpose_for_scores(mixed_query_layer)
+
+        past_key_value = (key_layer, value_layer)
+
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+        if (
+            self.position_embedding_type == "relative_key"
+            or self.position_embedding_type == "relative_key_query"
+        ):
+            seq_length = hidden_states.size()[1]
+            position_ids_l = torch.arange(
+                seq_length, dtype=torch.long, device=hidden_states.device
+            ).view(-1, 1)
+            position_ids_r = torch.arange(
+                seq_length, dtype=torch.long, device=hidden_states.device
+            ).view(1, -1)
+            distance = position_ids_l - position_ids_r
+            positional_embedding = self.distance_embedding(
+                distance + self.max_position_embeddings - 1
+            )
+            positional_embedding = positional_embedding.to(
+                dtype=query_layer.dtype
+            )  # fp16 compatibility
+
+            if self.position_embedding_type == "relative_key":
+                relative_position_scores = torch.einsum(
+                    "bhld,lrd->bhlr", query_layer, positional_embedding
+                )
+                attention_scores = attention_scores + relative_position_scores
+            elif self.position_embedding_type == "relative_key_query":
+                relative_position_scores_query = torch.einsum(
+                    "bhld,lrd->bhlr", query_layer, positional_embedding
+                )
+                relative_position_scores_key = torch.einsum(
+                    "bhrd,lrd->bhlr", key_layer, positional_embedding
+                )
+                attention_scores = (
+                    attention_scores
+                    + relative_position_scores_query
+                    + relative_position_scores_key
+                )
+
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+        if attention_mask is not None:
+            # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
+            attention_scores = attention_scores + attention_mask
+
+        # Normalize the attention scores to probabilities.
+        attention_probs = nn.Softmax(dim=-1)(attention_scores)
+
+        if is_cross_attention and self.save_attention:
+            self.save_attention_map(attention_probs)
+            attention_probs.register_hook(self.save_attn_gradients)
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+        attention_probs_dropped = self.dropout(attention_probs)
+
+        # Mask heads if we want to
+        if head_mask is not None:
+            attention_probs_dropped = attention_probs_dropped * head_mask
+
+        context_layer = torch.matmul(attention_probs_dropped, value_layer)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(*new_context_layer_shape)
+
+        outputs = (
+            (context_layer, attention_probs) if output_attentions else (context_layer,)
+        )
+
+        outputs = outputs + (past_key_value,)
+        return outputs
+
+
+class BertSelfOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class BertAttention(nn.Module):
+    def __init__(self, config, is_cross_attention=False):
+        super().__init__()
+        self.self = BertSelfAttention(config, is_cross_attention)
+        self.output = BertSelfOutput(config)
+        self.pruned_heads = set()
+
+    def prune_heads(self, heads):
+        if len(heads) == 0:
+            return
+        heads, index = find_pruneable_heads_and_indices(
+            heads,
+            self.self.num_attention_heads,
+            self.self.attention_head_size,
+            self.pruned_heads,
+        )
+
+        # Prune linear layers
+        self.self.query = prune_linear_layer(self.self.query, index)
+        self.self.key = prune_linear_layer(self.self.key, index)
+        self.self.value = prune_linear_layer(self.self.value, index)
+        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+
+        # Update hyper params and store pruned heads
+        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
+        self.self.all_head_size = (
+            self.self.attention_head_size * self.self.num_attention_heads
+        )
+        self.pruned_heads = self.pruned_heads.union(heads)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+    ):
+        self_outputs = self.self(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            encoder_hidden_states,
+            encoder_attention_mask,
+            past_key_value,
+            output_attentions,
+        )
+        attention_output = self.output(self_outputs[0], hidden_states)
+        outputs = (attention_output,) + self_outputs[
+            1:
+        ]  # add attentions if we output them
+        return outputs
+
+
+class BertIntermediate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        if isinstance(config.hidden_act, str):
+            self.intermediate_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.intermediate_act_fn = config.hidden_act
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.intermediate_act_fn(hidden_states)
+        return hidden_states
+
+
+class BertOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class BertLayer(nn.Module):
+    def __init__(self, config, layer_num):
+        super().__init__()
+        self.config = config
+        self.chunk_size_feed_forward = config.chunk_size_feed_forward
+        self.seq_len_dim = 1
+        self.attention = BertAttention(config)
+        self.layer_num = layer_num
+
+        # compatibility for ALBEF and BLIP
+        try:
+            # ALBEF & ALPRO
+            fusion_layer = self.config.fusion_layer
+            add_cross_attention = (
+                fusion_layer <= layer_num and self.config.add_cross_attention
+            )
+
+            self.fusion_layer = fusion_layer
+        except AttributeError:
+            # BLIP
+            self.fusion_layer = self.config.num_hidden_layers
+            add_cross_attention = self.config.add_cross_attention
+
+        # if self.config.add_cross_attention:
+        if add_cross_attention:
+            self.crossattention = BertAttention(
+                config, is_cross_attention=self.config.add_cross_attention
+            )
+        self.intermediate = BertIntermediate(config)
+        self.output = BertOutput(config)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+        mode=None,
+    ):
+        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
+        self_attn_past_key_value = (
+            past_key_value[:2] if past_key_value is not None else None
+        )
+        self_attention_outputs = self.attention(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            output_attentions=output_attentions,
+            past_key_value=self_attn_past_key_value,
+        )
+        attention_output = self_attention_outputs[0]
+
+        outputs = self_attention_outputs[1:-1]
+        present_key_value = self_attention_outputs[-1]
+
+        # TODO line 482 in albef/models/xbert.py
+        # compatibility for ALBEF and BLIP
+        if mode in ["multimodal", "fusion"] and hasattr(self, "crossattention"):
+            assert (
+                encoder_hidden_states is not None
+            ), "encoder_hidden_states must be given for cross-attention layers"
+
+            if isinstance(encoder_hidden_states, list):
+                cross_attention_outputs = self.crossattention(
+                    attention_output,
+                    attention_mask,
+                    head_mask,
+                    encoder_hidden_states[
+                        (self.layer_num - self.fusion_layer)
+                        % len(encoder_hidden_states)
+                    ],
+                    encoder_attention_mask[
+                        (self.layer_num - self.fusion_layer)
+                        % len(encoder_hidden_states)
+                    ],
+                    output_attentions=output_attentions,
+                )
+                attention_output = cross_attention_outputs[0]
+                outputs = outputs + cross_attention_outputs[1:-1]
+
+            else:
+                cross_attention_outputs = self.crossattention(
+                    attention_output,
+                    attention_mask,
+                    head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    output_attentions=output_attentions,
+                )
+                attention_output = cross_attention_outputs[0]
+                outputs = (
+                    outputs + cross_attention_outputs[1:-1]
+                )  # add cross attentions if we output attention weights
+        layer_output = apply_chunking_to_forward(
+            self.feed_forward_chunk,
+            self.chunk_size_feed_forward,
+            self.seq_len_dim,
+            attention_output,
+        )
+        outputs = (layer_output,) + outputs
+
+        outputs = outputs + (present_key_value,)
+
+        return outputs
+
+    def feed_forward_chunk(self, attention_output):
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
+        return layer_output
+
+
+class BertEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.layer = nn.ModuleList(
+            [BertLayer(config, i) for i in range(config.num_hidden_layers)]
+        )
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=False,
+        output_hidden_states=False,
+        return_dict=True,
+        mode="multimodal",
+    ):
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attentions = () if output_attentions else None
+        all_cross_attentions = (
+            () if output_attentions and self.config.add_cross_attention else None
+        )
+
+        next_decoder_cache = () if use_cache else None
+
+        try:
+            # ALBEF
+            fusion_layer = self.config.fusion_layer
+        except AttributeError:
+            # BLIP
+            fusion_layer = self.config.num_hidden_layers
+
+        if mode == "text":
+            start_layer = 0
+            # output_layer = self.config.fusion_layer
+            output_layer = fusion_layer
+
+        elif mode == "fusion":
+            # start_layer = self.config.fusion_layer
+            start_layer = fusion_layer
+            output_layer = self.config.num_hidden_layers
+
+        elif mode == "multimodal":
+            start_layer = 0
+            output_layer = self.config.num_hidden_layers
+
+        # compatibility for ALBEF and BLIP
+        # for i in range(self.config.num_hidden_layers):
+        for i in range(start_layer, output_layer):
+            layer_module = self.layer[i]
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+            past_key_value = past_key_values[i] if past_key_values is not None else None
+
+            # TODO pay attention to this.
+            if self.gradient_checkpointing and self.training:
+
+                if use_cache:
+                    logger.warn(
+                        "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                    )
+                    use_cache = False
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs, past_key_value, output_attentions)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(layer_module),
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    mode=mode,
+                )
+            else:
+                layer_outputs = layer_module(
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    past_key_value,
+                    output_attentions,
+                    mode=mode,
+                )
+
+            hidden_states = layer_outputs[0]
+            if use_cache:
+                next_decoder_cache += (layer_outputs[-1],)
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (layer_outputs[1],)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [
+                    hidden_states,
+                    next_decoder_cache,
+                    all_hidden_states,
+                    all_self_attentions,
+                    all_cross_attentions,
+                ]
+                if v is not None
+            )
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=next_decoder_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+            cross_attentions=all_cross_attentions,
+        )
+
+
+class BertPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        first_token_tensor = hidden_states[:, 0]
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class BertPredictionHeadTransform(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        if isinstance(config.hidden_act, str):
+            self.transform_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.transform_act_fn = config.hidden_act
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.transform_act_fn(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states)
+        return hidden_states
+
+
+class BertLMPredictionHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.transform = BertPredictionHeadTransform(config)
+
+        # The output weights are the same as the input embeddings, but there is
+        # an output-only bias for each token.
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+
+        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
+        self.decoder.bias = self.bias
+
+    def forward(self, hidden_states):
+        hidden_states = self.transform(hidden_states)
+        hidden_states = self.decoder(hidden_states)
+        return hidden_states
+
+
+class BertOnlyMLMHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.predictions = BertLMPredictionHead(config)
+
+    def forward(self, sequence_output):
+        prediction_scores = self.predictions(sequence_output)
+        return prediction_scores
+
+
+class BertPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = BertConfig
+    base_model_prefix = "bert"
+    _keys_to_ignore_on_load_missing = [r"position_ids"]
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+
+class BertModel(BertPreTrainedModel):
+    """
+    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
+    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
+    all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
+    Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
+    argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an
+    input to the forward pass.
+    """
+
+    def __init__(self, config, add_pooling_layer=True):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = BertEmbeddings(config)
+
+        self.encoder = BertEncoder(config)
+
+        self.pooler = BertPooler(config) if add_pooling_layer else None
+
+        self.init_weights()
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    def _prune_heads(self, heads_to_prune):
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    def get_extended_attention_mask(
+        self,
+        attention_mask: Tensor,
+        input_shape: Tuple[int],
+        device: device,
+        is_decoder: bool,
+    ) -> Tensor:
+        """
+        Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
+
+        Arguments:
+            attention_mask (:obj:`torch.Tensor`):
+                Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
+            input_shape (:obj:`Tuple[int]`):
+                The shape of the input to the model.
+            device: (:obj:`torch.device`):
+                The device of the input to the model.
+
+        Returns:
+            :obj:`torch.Tensor` The extended attention mask, with a the same dtype as :obj:`attention_mask.dtype`.
+        """
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        if attention_mask.dim() == 3:
+            extended_attention_mask = attention_mask[:, None, :, :]
+        elif attention_mask.dim() == 2:
+            # Provided a padding mask of dimensions [batch_size, seq_length]
+            # - if the model is a decoder, apply a causal mask in addition to the padding mask
+            # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
+            if is_decoder:
+                batch_size, seq_length = input_shape
+
+                seq_ids = torch.arange(seq_length, device=device)
+                causal_mask = (
+                    seq_ids[None, None, :].repeat(batch_size, seq_length, 1)
+                    <= seq_ids[None, :, None]
+                )
+                # in case past_key_values are used we need to add a prefix ones mask to the causal mask
+                # causal and attention masks must have same type with pytorch version < 1.3
+                causal_mask = causal_mask.to(attention_mask.dtype)
+
+                if causal_mask.shape[1] < attention_mask.shape[1]:
+                    prefix_seq_len = attention_mask.shape[1] - causal_mask.shape[1]
+                    causal_mask = torch.cat(
+                        [
+                            torch.ones(
+                                (batch_size, seq_length, prefix_seq_len),
+                                device=device,
+                                dtype=causal_mask.dtype,
+                            ),
+                            causal_mask,
+                        ],
+                        axis=-1,
+                    )
+
+                extended_attention_mask = (
+                    causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
+                )
+            else:
+                extended_attention_mask = attention_mask[:, None, None, :]
+        else:
+            raise ValueError(
+                "Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
+                    input_shape, attention_mask.shape
+                )
+            )
+
+        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
+        # masked positions, this operation will create a tensor which is 0.0 for
+        # positions we want to attend and -10000.0 for masked positions.
+        # Since we are adding it to the raw scores before the softmax, this is
+        # effectively the same as removing these entirely.
+        extended_attention_mask = extended_attention_mask.to(
+            dtype=self.dtype
+        )  # fp16 compatibility
+        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+        return extended_attention_mask
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        is_decoder=False,
+        mode="multimodal",
+    ):
+        r"""
+        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
+            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+            the model is configured as a decoder.
+        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
+            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
+            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
+            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
+        use_cache (:obj:`bool`, `optional`):
+            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
+            decoding (see :obj:`past_key_values`).
+        """
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        if is_decoder:
+            use_cache = use_cache if use_cache is not None else self.config.use_cache
+        else:
+            use_cache = False
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time"
+            )
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+            batch_size, seq_length = input_shape
+            device = input_ids.device
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+            batch_size, seq_length = input_shape
+            device = inputs_embeds.device
+        elif encoder_embeds is not None:
+            input_shape = encoder_embeds.size()[:-1]
+            batch_size, seq_length = input_shape
+            device = encoder_embeds.device
+        else:
+            raise ValueError(
+                "You have to specify either input_ids or inputs_embeds or encoder_embeds"
+            )
+
+        # past_key_values_length
+        past_key_values_length = (
+            past_key_values[0][0].shape[2] if past_key_values is not None else 0
+        )
+
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                ((batch_size, seq_length + past_key_values_length)), device=device
+            )
+
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(
+            attention_mask, input_shape, device, is_decoder
+        )
+
+        # If a 2D or 3D attention mask is provided for the cross-attention
+        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+        if encoder_hidden_states is not None:
+            if type(encoder_hidden_states) == list:
+                encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states[
+                    0
+                ].size()
+            else:
+                (
+                    encoder_batch_size,
+                    encoder_sequence_length,
+                    _,
+                ) = encoder_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+
+            if type(encoder_attention_mask) == list:
+                encoder_extended_attention_mask = [
+                    self.invert_attention_mask(mask) for mask in encoder_attention_mask
+                ]
+            elif encoder_attention_mask is None:
+                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
+                encoder_extended_attention_mask = self.invert_attention_mask(
+                    encoder_attention_mask
+                )
+            else:
+                encoder_extended_attention_mask = self.invert_attention_mask(
+                    encoder_attention_mask
+                )
+        else:
+            encoder_extended_attention_mask = None
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+
+        if encoder_embeds is None:
+            embedding_output = self.embeddings(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                token_type_ids=token_type_ids,
+                inputs_embeds=inputs_embeds,
+                past_key_values_length=past_key_values_length,
+            )
+        else:
+            embedding_output = encoder_embeds
+
+        encoder_outputs = self.encoder(
+            embedding_output,
+            attention_mask=extended_attention_mask,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_extended_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            mode=mode,
+        )
+        sequence_output = encoder_outputs[0]
+        pooled_output = (
+            self.pooler(sequence_output) if self.pooler is not None else None
+        )
+
+        if not return_dict:
+            return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPoolingAndCrossAttentions(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            past_key_values=encoder_outputs.past_key_values,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+            cross_attentions=encoder_outputs.cross_attentions,
+        )
+
+
+class BertForMaskedLM(BertPreTrainedModel):
+
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+    _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.bert = BertModel(config, add_pooling_layer=False)
+        self.cls = BertOnlyMLMHead(config)
+
+        self.init_weights()
+
+    def get_output_embeddings(self):
+        return self.cls.predictions.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.cls.predictions.decoder = new_embeddings
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        # token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        is_decoder=False,
+        mode="multimodal",
+        soft_labels=None,
+        alpha=0,
+        return_logits=False,
+    ):
+        r"""
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Labels for computing the masked language modeling loss. Indices should be in ``[-100, 0, ...,
+            config.vocab_size]`` (see ``input_ids`` docstring) Tokens with indices set to ``-100`` are ignored
+            (masked), the loss is only computed for the tokens with labels in ``[0, ..., config.vocab_size]``
+        """
+
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            # token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_embeds=encoder_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            is_decoder=is_decoder,
+            mode=mode,
+        )
+
+        sequence_output = outputs[0]
+        prediction_scores = self.cls(sequence_output)
+
+        if return_logits:
+            return prediction_scores
+
+        masked_lm_loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()  # -100 index = padding token
+            masked_lm_loss = loss_fct(
+                prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)
+            )
+
+        if soft_labels is not None:
+            loss_distill = -torch.sum(
+                F.log_softmax(prediction_scores, dim=-1) * soft_labels, dim=-1
+            )
+            loss_distill = loss_distill[labels != -100].mean()
+            masked_lm_loss = (1 - alpha) * masked_lm_loss + alpha * loss_distill
+
+        if not return_dict:
+            output = (prediction_scores,) + outputs[2:]
+            return (
+                ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+            )
+
+        return MaskedLMOutput(
+            loss=masked_lm_loss,
+            logits=prediction_scores,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self, input_ids, attention_mask=None, **model_kwargs
+    ):
+        input_shape = input_ids.shape
+        effective_batch_size = input_shape[0]
+
+        #  add a dummy token
+        assert (
+            self.config.pad_token_id is not None
+        ), "The PAD token should be defined for generation"
+        attention_mask = torch.cat(
+            [attention_mask, attention_mask.new_zeros((attention_mask.shape[0], 1))],
+            dim=-1,
+        )
+        dummy_token = torch.full(
+            (effective_batch_size, 1),
+            self.config.pad_token_id,
+            dtype=torch.long,
+            device=input_ids.device,
+        )
+        input_ids = torch.cat([input_ids, dummy_token], dim=1)
+
+        return {"input_ids": input_ids, "attention_mask": attention_mask}
+
+
+class BertLMHeadModel(BertPreTrainedModel):
+
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+    _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.bert = BertModel(config, add_pooling_layer=False)
+        self.cls = BertOnlyMLMHead(config)
+
+        self.init_weights()
+
+    def get_output_embeddings(self):
+        return self.cls.predictions.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.cls.predictions.decoder = new_embeddings
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        labels=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        return_logits=False,
+        is_decoder=True,
+        reduction="mean",
+        mode="multimodal",
+        soft_labels=None,
+        alpha=0,
+    ):
+        r"""
+        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
+            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+            the model is configured as a decoder.
+        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
+            ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring) Tokens with indices set to ``-100`` are
+            ignored (masked), the loss is only computed for the tokens with labels n ``[0, ..., config.vocab_size]``
+        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
+            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
+            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
+            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
+        use_cache (:obj:`bool`, `optional`):
+            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
+            decoding (see :obj:`past_key_values`).
+        Returns:
+        Example::
+            >>> from transformers import BertTokenizer, BertLMHeadModel, BertConfig
+            >>> import torch
+            >>> tokenizer = BertTokenizer.from_pretrained('bert-base-cased')
+            >>> config = BertConfig.from_pretrained("bert-base-cased")
+            >>> model = BertLMHeadModel.from_pretrained('bert-base-cased', config=config)
+            >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+            >>> outputs = model(**inputs)
+            >>> prediction_logits = outputs.logits
+        """
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+        if labels is not None:
+            use_cache = False
+
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            is_decoder=is_decoder,
+            mode=mode,
+        )
+
+        sequence_output = outputs[0]
+        prediction_scores = self.cls(sequence_output)
+
+        if return_logits:
+            return prediction_scores[:, :-1, :].contiguous()
+
+        lm_loss = None
+        if labels is not None:
+            # we are doing next-token prediction; shift prediction scores and input ids by one
+            shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous()
+            labels = labels[:, 1:].contiguous()
+            loss_fct = CrossEntropyLoss(reduction=reduction, label_smoothing=0.1)
+            lm_loss = loss_fct(
+                shifted_prediction_scores.view(-1, self.config.vocab_size),
+                labels.view(-1),
+            )
+            if reduction == "none":
+                lm_loss = lm_loss.view(prediction_scores.size(0), -1).sum(1)
+
+        if soft_labels is not None:
+            loss_distill = -torch.sum(
+                F.log_softmax(shifted_prediction_scores, dim=-1) * soft_labels, dim=-1
+            )
+            loss_distill = (loss_distill * (labels != -100)).sum(1)
+            lm_loss = (1 - alpha) * lm_loss + alpha * loss_distill
+
+        if not return_dict:
+            output = (prediction_scores,) + outputs[2:]
+            return ((lm_loss,) + output) if lm_loss is not None else output
+
+        # for embedding of the desc in the decoder, return the last_hidden_state (embeds)
+        if (type(self).__name__ == 'XBertLMHeadDecoder') and (outputs.last_hidden_state.shape[1] == 200):
+            return [CausalLMOutputWithCrossAttentions(
+                loss=lm_loss,
+                logits=prediction_scores,
+                past_key_values=outputs.past_key_values,
+                hidden_states=outputs.hidden_states,
+                attentions=outputs.attentions,
+                cross_attentions=outputs.cross_attentions,
+            ), outputs.last_hidden_state]
+        else:
+            return CausalLMOutputWithCrossAttentions(
+                loss=lm_loss,
+                logits=prediction_scores,
+                past_key_values=outputs.past_key_values,
+                hidden_states=outputs.hidden_states,
+                attentions=outputs.attentions,
+                cross_attentions=outputs.cross_attentions,
+            )
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past=None, attention_mask=None, **model_kwargs
+    ):
+        input_shape = input_ids.shape
+        # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
+        if attention_mask is None:
+            attention_mask = input_ids.new_ones(input_shape)
+
+        # cut decoder_input_ids if past is used
+        if past is not None:
+            input_ids = input_ids[:, -1:]
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "past_key_values": past,
+            "encoder_hidden_states": model_kwargs.get("encoder_hidden_states", None),
+            "encoder_attention_mask": model_kwargs.get("encoder_attention_mask", None),
+            "is_decoder": True,
+        }
+
+    def _reorder_cache(self, past, beam_idx):
+        reordered_past = ()
+        for layer_past in past:
+            reordered_past += (
+                tuple(
+                    past_state.index_select(0, beam_idx) for past_state in layer_past
+                ),
+            )
+        return reordered_past
+
+
+class XBertLMHeadDecoder(BertLMHeadModel):
+    """
+    This class decouples the decoder forward logic from the VL model.
+    In this way, different VL models can share this decoder as long as
+    they feed encoder_embeds as required.
+    """
+
+    @classmethod
+    def from_config(cls, cfg, from_pretrained=False):
+        med_config_path = get_abs_path(cfg.get("med_config_path"))
+        med_config = BertConfig.from_json_file(med_config_path)
+
+        if from_pretrained:
+            return cls.from_pretrained("bert-base-uncased", config=med_config)
+        else:
+            return cls(config=med_config)
+
+    def generate_from_encoder(
+        self,
+        tokenized_prompt,
+        visual_embeds,
+        sep_token_id,
+        pad_token_id,
+        use_nucleus_sampling=False,
+        num_beams=3,
+        max_length=30,
+        min_length=10,
+        top_p=0.9,
+        repetition_penalty=1.0,
+        **kwargs
+    ):
+
+        if not use_nucleus_sampling:
+            num_beams = num_beams
+            visual_embeds = visual_embeds.repeat_interleave(num_beams, dim=0)
+
+        image_atts = torch.ones(visual_embeds.size()[:-1], dtype=torch.long).to(
+            self.device
+        )
+
+        model_kwargs = {
+            "encoder_hidden_states": visual_embeds,
+            "encoder_attention_mask": image_atts,
+        }
+
+        if use_nucleus_sampling:
+            # nucleus sampling
+            outputs = self.generate(
+                input_ids=tokenized_prompt.input_ids,
+                max_length=max_length,
+                min_length=min_length,
+                do_sample=True,
+                top_p=top_p,
+                num_return_sequences=1,
+                eos_token_id=sep_token_id,
+                pad_token_id=pad_token_id,
+                repetition_penalty=1.1,
+                **model_kwargs
+            )
+        else:
+            # beam search
+            outputs = self.generate(
+                input_ids=tokenized_prompt.input_ids,
+                max_length=max_length,
+                min_length=min_length,
+                num_beams=num_beams,
+                eos_token_id=sep_token_id,
+                pad_token_id=pad_token_id,
+                repetition_penalty=repetition_penalty,
+                **model_kwargs
+            )
+
+        return outputs
+
+
+class XBertEncoder(BertModel, BaseEncoder):
+    def __init__(self, config, add_pooling_layer):
+        super().__init__(config, add_pooling_layer)
+
+        self.cls = BertOnlyMLMHead(config)
+        self.init_weights()
+    
+    def get_output_embeddings(self):
+        return self.cls.predictions.decoder
+    
+    @classmethod
+    def from_config(cls, cfg, from_pretrained=False):
+        configs_root = os.environ.get("CONFIGS_ROOT", "../ckpt")
+        med_config_path = os.path.join(configs_root, "bert-base-chinese/config.json")
+        med_config = BertConfig.from_json_file(med_config_path)
+
+        if from_pretrained:
+            model = cls.from_pretrained(
+                os.path.join(configs_root, "bert-base-chinese"),
+                config=med_config,
+                add_pooling_layer=False
+            )
+        else:
+            model = cls(config=med_config, add_pooling_layer=False)
+        
+        for name, param in model.named_parameters():
+            if 'cls.predictions' in name:
+                param.requires_grad = False
+        
+        return model
+
+    def forward_automask(self, tokenized_text, visual_embeds, **kwargs):
+        image_atts = torch.ones(visual_embeds.size()[:-1], dtype=torch.long).to(
+            self.device
+        )
+
+        text = tokenized_text
+        text_output = super().forward(
+            text.input_ids,
+            attention_mask=text.attention_mask,
+            encoder_hidden_states=visual_embeds,
+            encoder_attention_mask=image_atts,
+            return_dict=True,
+        )
+
+        return text_output
+
+    def forward_text(self, tokenized_text, **kwargs):
+        text = tokenized_text
+        token_type_ids = kwargs.get("token_type_ids", None)
+
+        text_output = super().forward(
+            text.input_ids,
+            attention_mask=text.attention_mask,
+            token_type_ids=token_type_ids,
+            return_dict=True,
+            mode="text",
+        )
+
+        return text_output
+    
+    def forward_mlm(self, 
+            input_ids, 
+            attention_mask,
+            encoder_hidden_states,
+            encoder_attention_mask,
+            return_dict,
+            labels
+            ):
+
+        # slef.bert()
+        outputs = super().forward(
+            input_ids=input_ids, 
+            attention_mask=attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            return_dict=return_dict
+        )
+
+        sequence_output = outputs[0]
+        prediction_scores = self.cls(sequence_output)
+
+        # if return_logits:
+        #     return prediction_scores
+
+        masked_lm_loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()  # -100 index = padding token
+            masked_lm_loss = loss_fct(
+                prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)
+            )
+
+        if not return_dict:
+            output = (prediction_scores,) + outputs[2:]
+            return (
+                ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+            )
+
+        return MaskedLMOutput(
+            loss=masked_lm_loss,
+            logits=prediction_scores,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

RADAR_inference/dynamic_network_architectures/vision_branch.py

@@ -0,0 +1,160 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+import torch.nn as nn
+import sys
+from monai.utils import deprecated_arg
+import pydoc
+import warnings
+from typing import Union
+import torch.nn.functional as F
+
+
+class VisionBranch(nn.Module):
+    def __init__(
+        self,
+        in_channels=1
+    ) -> None:
+        super().__init__()
+
+        # define LightDecoderUNet
+        self.UNet = self.get_network_from_plans(
+            arch_class_name="dynamic_network_architectures.architectures.unet_lightdecoder.PlainConvUNetLightD",
+            arch_kwargs={
+                "n_stages": 6,
+                "features_per_stage": [32, 64, 128, 256, 320, 320],
+                "conv_op": "torch.nn.modules.conv.Conv3d",
+                "kernel_sizes": [[1, 3, 3], [1, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3]],
+                "strides": [[1, 1, 1], [1, 2, 2], [1, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]],
+                "n_conv_per_stage": [2, 2, 2, 2, 2, 2],
+                "n_conv_per_stage_decoder": [1, 1, 1, 1, 1],
+                "conv_bias": True,
+                "norm_op": "torch.nn.BatchNorm3d",
+                "norm_op_kwargs": {},
+                "dropout_op": None,
+                "dropout_op_kwargs": None,
+                "nonlin": "torch.nn.ReLU",
+                "nonlin_kwargs": {"inplace": True},
+            },
+            arch_kwargs_req_import=["conv_op", "norm_op", "dropout_op", "nonlin"],
+            input_channels=1,
+            output_channels=37,
+            allow_init=True,
+            deep_supervision=True,
+        )
+        
+        self.proj1 = nn.Conv3d(320, 256, kernel_size=1)
+        self.proj2 = nn.Conv3d(320, 256, kernel_size=1)
+        self.proj3 = nn.Conv3d(256, 256, kernel_size=1)
+        
+        self.organs = [
+            '肾上腺', '主动脉', '竖脊肌', '脑', '锁骨', '大肠', '十二指肠', '食管', '面部', '股骨', 
+            '胆囊', "臀肌", '心脏', '髋关节', '肱骨', '髂动脉', '髂静脉', '髂腰肌', '下腔静脉', '肾', 
+            '肝', '肺', '胰腺', '门静脉', '肺动脉', '肋骨', '骶骨', '肩胛骨', '小肠', '脾', 
+            '胃', '气管', '膀胱', '颈椎', '腰椎', '胸椎'
+        ]
+        # organ_dict = {
+        #     "肾上腺": "adrenal gland", "主动脉": "aorta", "竖脊肌": "erector spinae muscle", "脑": "brain", "锁骨": "clavicle", "大肠": "large bowel", "十二指肠": "duodenum", 
+        #     "食管": "esophagus", "面部": "face", "股骨": "femur", "胆囊": "gallbladder", "臀肌": "gluteus muscle", "心脏": "heart", "髋关节": "hip joint", "肱骨": "humerus", 
+        #     "髂动脉": "iliac artery", "髂静脉": "iliac vena", "髂腰肌": "iliopsoas muscle", "下腔静脉": "inferior vena cava", "肾": "kidney", "肝": "liver", "肺": "lung", 
+        #     "胰腺": "pancreas", "门静脉": "portal vein", "肺动脉": "pulmonary artery", "肋骨": "rib", "骶骨": "sacrum", "肩胛骨": "scapula", "小肠": "small bowel", "脾": "spleen", 
+        #     "胃": "stomach", "气管": "trachea", "膀胱": "bladder", "颈椎": "cervical vertebrae", "腰椎": "lumbar vertebrae", "胸椎": "thoracic vertebrae"
+        # }
+
+    def get_network_from_plans(sefl, arch_class_name, arch_kwargs, arch_kwargs_req_import, input_channels, output_channels,
+                           allow_init=True, deep_supervision: Union[bool, None] = None):
+        network_class = arch_class_name
+        architecture_kwargs = dict(**arch_kwargs)
+        for ri in arch_kwargs_req_import:
+            if architecture_kwargs[ri] is not None:
+                architecture_kwargs[ri] = pydoc.locate(architecture_kwargs[ri])
+
+        nw_class = pydoc.locate(network_class)
+
+        if deep_supervision is not None:
+            architecture_kwargs['deep_supervision'] = deep_supervision
+
+        network = nw_class(
+            input_channels=input_channels,
+            num_classes=output_channels,
+            **architecture_kwargs
+        )
+
+        if hasattr(network, 'initialize') and allow_init:
+            network.apply(network.initialize)
+
+        return network
+
+    def forward(self, x, y):
+        skips, segs = self.UNet(x)
+        
+        scale1 = self.proj1(skips[-1])
+        scale2 = self.proj2(skips[-2])
+        scale3 = self.proj3(skips[-3])
+        
+        # process mask
+        pred_logit = segs[0]
+        seg_probs = torch.softmax(pred_logit, 1)
+        target_size = [pred_logit.shape[-3], pred_logit.shape[-2]*2, pred_logit.shape[-1]*2]
+        pred_logit = F.interpolate(pred_logit, size=target_size, mode='trilinear', align_corners=False)
+        pred_mask = torch.softmax(pred_logit, 1)
+        pred_mask = pred_mask.argmax(1)
+        y = pred_mask
+        
+        res_x1 = scale1.flatten(2).transpose(1, 2)
+        res_x2 = scale2.flatten(2).transpose(1, 2)
+        res_x3 = scale3.flatten(2).transpose(1, 2)
+        
+        B, L1, _ = res_x1.size()
+        B, L2, _ = res_x2.size()
+        B, L3, _ = res_x3.size()
+        
+        with torch.no_grad():
+            organ_token_flags1 = torch.zeros(B, len(self.organs), L1, dtype=bool).to(x.device)
+            organ_token_flags2 = torch.zeros(B, len(self.organs), L2, dtype=bool).to(x.device)
+            organ_token_flags3 = torch.zeros(B, len(self.organs), L3, dtype=bool).to(x.device)
+            
+            b = x.size(0)
+            for i in range(b):
+                unique_values = torch.unique(y[i])
+                unique_values = unique_values[unique_values != 0]
+                if unique_values.tolist() == []:
+                    continue
+                masks = torch.stack([torch.eq(y[i], uv) for uv in unique_values]).float()
+                
+                highlight_tokens3 = F.max_pool3d(
+                    masks.unsqueeze(1),
+                    kernel_size=(2, 8, 8),
+                    stride=(2, 8, 8)
+                ).flatten(1) > 0
+
+                highlight_tokens2 = F.max_pool3d(
+                    masks.unsqueeze(1),
+                    kernel_size=(4, 16, 16),
+                    stride=(4, 16, 16)
+                ).flatten(1) > 0
+
+                highlight_tokens1 = F.max_pool3d(
+                    masks.unsqueeze(1),
+                    kernel_size=(8, 32, 32),
+                    stride=(8, 32, 32)
+                ).flatten(1) > 0
+
+                organ_token_flags1[i][unique_values.long() - 1] = highlight_tokens1 > 0
+                organ_token_flags2[i][unique_values.long() - 1] = highlight_tokens2 > 0
+                organ_token_flags3[i][unique_values.long() - 1] = highlight_tokens3 > 0
+
+        return seg_probs, pred_mask, res_x1, res_x2, res_x3, organ_token_flags1, organ_token_flags2, organ_token_flags3

RADAR_inference/inference_demo.py

@@ -0,0 +1,630 @@
+import os
+import re
+import json
+import argparse
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+from torch.utils.data import Dataset, DataLoader, DistributedSampler
+from torch.utils.data.dataloader import default_collate
+from monai import transforms
+from monai.data.utils import dense_patch_slices
+from typing import Any, Callable, List, Sequence, Tuple, Union
+import datetime
+import SimpleITK as sitk
+import torch.nn.functional as F
+from pathlib import Path
+from copy import deepcopy
+import torch.distributed as dist
+from dynamic_network_architectures.med import XBertEncoder, XBertLMHeadDecoder
+from dynamic_network_architectures.vision_branch import VisionBranch
+from transformers import BertTokenizer
+
+
+model_root = os.environ.get("MODEL_ROOT", "../ckpt")
+configs_root = os.environ.get("CONFIGS_ROOT", "../ckpt")
+
+def masks_to_boxes_3d(masks):
+    """Compute the bounding boxes around the provided 3D masks
+
+    The masks should be in format [N, D, H, W] where N is the number of masks, (D, H, W) are the spatial dimensions.
+
+    Returns a [N, 6] tensor, with the boxes in min_x, min_y, min_z, max_x, max_y, max_z format
+    """
+    if masks.numel() == 0:
+        return torch.zeros((0, 6), device=masks.device)
+
+    d, h, w = masks.shape[-3:]
+
+    z = torch.arange(0, d, dtype=torch.float, device=masks.device)
+    y = torch.arange(0, h, dtype=torch.float, device=masks.device)
+    x = torch.arange(0, w, dtype=torch.float, device=masks.device)
+
+    z, y, x = torch.meshgrid(z, y, x, indexing='ij')
+
+    x_mask = (masks * x.unsqueeze(0))
+    x_max = x_mask.flatten(1).max(-1).values
+    x_min = x_mask.masked_fill(~masks.bool(), float('inf')).flatten(1).min(-1).values
+
+    y_mask = (masks * y.unsqueeze(0))
+    y_max = y_mask.flatten(1).max(-1).values
+    y_min = y_mask.masked_fill(~masks.bool(), float('inf')).flatten(1).min(-1).values
+
+    z_mask = (masks * z.unsqueeze(0))
+    z_max = z_mask.flatten(1).max(-1).values
+    z_min = z_mask.masked_fill(~masks.bool(), float('inf')).flatten(1).min(-1).values
+
+    return torch.stack([x_min, y_min, z_min, x_max, y_max, z_max], dim=1)
+
+def collate_fn(batch):
+    return batch[0]
+
+@torch.no_grad()
+def all_gather(data):
+    world_size = dist.get_world_size()
+    if world_size == 1:
+        return [data]
+    data_list = [None] * world_size
+    dist.all_gather_object(data_list, data)
+    return data_list
+
+def _get_scan_interval(
+    image_size: Sequence[int], roi_size: Sequence[int], num_spatial_dims: int, overlap: float
+) -> Tuple[int, ...]:
+        """
+        Compute scan interval according to the image size, roi size and overlap.
+        Scan interval will be `int((1 - overlap) * roi_size)`, if interval is 0,
+        use 1 instead to make sure sliding window works.
+
+        """
+        if len(image_size) != num_spatial_dims:
+            raise ValueError("image coord different from spatial dims.")
+        if len(roi_size) != num_spatial_dims:
+            raise ValueError("roi coord different from spatial dims.")
+
+        scan_interval = []
+        for i in range(num_spatial_dims):
+            if roi_size[i] == image_size[i]:
+                scan_interval.append(int(roi_size[i]))
+            else:
+                interval = int(roi_size[i] * (1 - overlap))
+                scan_interval.append(interval if interval > 0 else 1)
+        return tuple(scan_interval)
+
+def center_crop(image, mask, crop_size):
+    x_min, y_min, z_min, x_max, y_max, z_max = masks_to_boxes_3d(mask)[0].long()
+    
+    crop_d, crop_h, crop_w = max(crop_size[0], z_max - z_min), max(crop_size[1], y_max - y_min), max(crop_size[2], x_max - x_min)
+
+    cx = (x_min + x_max) // 2
+    cy = (y_min + y_max) // 2
+    cz = (z_min + z_max) // 2
+    
+    d, h, w = image.shape[-3:]
+
+    x_start = max(0, cx - crop_w // 2)
+    x_end = min(w, x_start + crop_w)
+    if x_end - x_start < crop_w:
+        x_start = max(0, x_end - crop_w)
+    
+    y_start = max(0, cy - crop_h // 2)
+    y_end = min(h, y_start + crop_h)
+    if y_end - y_start < crop_h:
+        y_start = max(0, y_end - crop_h)
+    
+    z_start = max(0, cz - crop_d // 2)
+    z_end = min(d, z_start + crop_d)
+    if z_end - z_start < crop_d:
+        z_start = max(0, z_end - crop_d)
+    
+    return image[..., z_start:z_end, y_start:y_end, x_start:x_end], mask[..., z_start:z_end, y_start:y_end, x_start:x_end]
+
+class DataFolder(Dataset):
+    def __init__(self, img_dir):
+        super().__init__()
+        
+        patient_list = os.listdir(img_dir)
+        self.img_paths = [
+            os.path.join(img_dir, p)
+            for p in patient_list
+        ]
+        
+        self.pad_func = transforms.SpatialPadd(
+            keys=["image"], 
+            spatial_size=(96, 256, 384), 
+            mode='constant', 
+            constant_values=0,
+            method="end"
+        )
+
+        self.organs = [
+            '肾上腺', '主动脉', '竖脊肌', '脑', '锁骨', '大肠', '十二指肠', '食管', '面部', '股骨', 
+            '胆囊', "臀肌", '心脏', '髋关节', '肱骨', '髂动脉', '髂静脉', '髂腰肌', '下腔静脉', '肾', 
+            '肝', '肺', '胰腺', '门静脉', '肺动脉', '肋骨', '骶骨', '肩胛骨', '小肠', '脾', 
+            '胃', '气管', '膀胱', '颈椎', '腰椎', '胸椎'
+        ]
+        # organ_dict = {
+        #     "肾上腺": "adrenal gland", "主动脉": "aorta", "竖脊肌": "erector spinae muscle", "脑": "brain", "锁骨": "clavicle", "大肠": "large bowel", "十二指肠": "duodenum", 
+        #     "食管": "esophagus", "面部": "face", "股骨": "femur", "胆囊": "gallbladder", "臀肌": "gluteus muscle", "心脏": "heart", "髋关节": "hip joint", "肱骨": "humerus", 
+        #     "髂动脉": "iliac artery", "髂静脉": "iliac vena", "髂腰肌": "iliopsoas muscle", "下腔静脉": "inferior vena cava", "肾": "kidney", "肝": "liver", "肺": "lung", 
+        #     "胰腺": "pancreas", "门静脉": "portal vein", "肺动脉": "pulmonary artery", "肋骨": "rib", "骶骨": "sacrum", "肩胛骨": "scapula", "小肠": "small bowel", "脾": "spleen", 
+        #     "胃": "stomach", "气管": "trachea", "膀胱": "bladder", "颈椎": "cervical vertebrae", "腰椎": "lumbar vertebrae", "胸椎": "thoracic vertebrae"
+        # }
+        
+        self.test_items = ['主动脉_主动脉夹层', '主动脉_主动脉瘤', '主动脉_粥样硬化', '主动脉_钙化', '十二指肠_占位', '十二指肠_囊袋状突出影', '十二指肠_憩室', '十二指肠_梗阻', '十二指肠_溃疡', '大肠_克罗恩病', '大肠_大肠（壁）钙化', '大肠_急慢性（结）肠炎', '大肠_浆膜面毛糙', '大肠_溃疡性结肠炎', '大肠_直肠癌', '大肠_积液积气', '大肠_结肠癌', '大肠_肠壁毛糙', '大肠_肠壁水肿', '大肠_肠套叠', '大肠_肠憩室', '大肠_肠梗阻', '大肠_肠穿孔', '大肠_肠道扩张', '大肠_脂肪间隙模糊', '大肠_阑尾炎', '大肠_阑尾粪石', '小肠_克罗恩病', '小肠_套叠', '小肠_扭转', '小肠_梗阻', '小肠_淋巴瘤', '小肠_积气积液', '小肠_系膜指膜炎', '小肠_系膜淋巴结肿大', '小肠_肠壁增厚', '小肠_肠管扩张', '小肠_脂肪瘤', '小肠_间质瘤（胃肠间质瘤-gist）', '小肠_（急慢性）小肠炎', '心脏_心包积液', '心脏_心影（脏）增大', '肋骨_转移瘤（乳腺癌 骨转移）', '肋骨_骨折', '肋骨_骨质破坏', '肝_低密度影', '肝_格林森鞘积液', '肝_比例失调', '肝_波浪状改变', '肝_硬化', '肝_结节状强化', '肝_肝内胆管扩张', '肝_肝内胆管结石', '肝_肝内钙化灶', '肝_肝囊肿', '肝_肝细胞癌', '肝_肝胆管内高密度影', '肝_肝血管瘤', '肝_胆管癌', '肝_脂肪肝', '肝_脓肿', '肝_转移瘤', '肝_边缘不规则', '肺_斑片影', '肺_气胸', '肺_结节', '肺_肺占位', '肺_肺萎陷', '肺_胸腔积液', '肺_膨胀不全', '肺_转移瘤', '肺_钙化灶', '肺_高密度影', '肾_低密度影', '肾_囊肿', '肾_多囊肾', '肾_实质变薄', '肾_无强化囊性灶', '肾_肾动脉瘤', '肾_肾盂扩张', '肾_肾盂癌', '肾_肾盂积水', '肾_肾细胞癌（透明细胞癌）', '肾_肾萎缩', '肾_肾血管平滑肌脂肪瘤', '肾_肾（盂）结石', '肾_高密度影', '肾上腺_增生', '肾上腺_结节', '肾上腺_脂肪瘤', '肾上腺_腺瘤', '肾上腺_转移瘤', '肾上腺_钙化', '胃_壁水肿', '胃_扩张', '胃_胃底静脉曲张', '胃_胃溃疡', '胃_胃癌', '胃_间质瘤（gist）', '胆囊_结石', '胆囊_结节状致密影', '胆囊_胆囊增大', '胆囊_胆囊炎', '胆囊_胆囊癌', '胆囊_胆囊腺肌症', '胆囊_胆管壁增厚', '胆囊_胆管扩张', '胆囊_胆管炎', '胆囊_胆管癌', '胆囊_胆管积气', '胆囊_胆管结石', '胆囊_高密度影', '胆囊_黄色肉芽肿', '胰腺_低密度影', '胰腺_囊肿', '胰腺_围脂肪间隙模糊', '胰腺_肿瘤或胰腺癌', '胰腺_胰周假性囊肿', '胰腺_胰管扩张', '胰腺_胰管结石', '胰腺_胰腺炎', '胰腺_胰腺饱满', '胰腺_萎缩', '脾_低密度灶', '脾_副脾', '脾_囊肿', '脾_梗死', '脾_片状低密度区', '脾_脾大', '脾_脾脏淋巴瘤', '脾_钙化', '膀胱_憩室', '膀胱_结石', '膀胱_膀胱壁毛糙', '膀胱_膀胱炎', '膀胱_膀胱癌', '膀胱_软组织密度影', '门静脉_增宽', '门静脉_栓塞', '门静脉_高压', '食管_增粗迂曲血管影', '食管_管壁增厚', '食管_裂孔疝', '食管_静脉扩张迂曲', '食管_静脉曲张', '骶骨_骨炎']
+        self.english_mapping = {'主动脉_主动脉夹层': 'Aorta_Aortic dissection', '主动脉_主动脉瘤': 'Aorta_Aortic aneurysm', '主动脉_粥样硬化': 'Aorta_Atherosclerosis', '主动脉_钙化': 'Aorta_Calcification', '十二指肠_占位': 'Duodenum_Mass', '十二指肠_囊袋状突出影': 'Duodenum_Saccular outpouching', '十二指肠_憩室': 'Duodenum_Diverticulum', '十二指肠_梗阻': 'Duodenum_Obstruction', '十二指肠_溃疡': 'Duodenum_Ulcer', '大肠_克罗恩病': "Large bowel_Crohn's disease", '大肠_大肠（壁）钙化': 'Large bowel_Mural calcification', '大肠_急慢性（结）肠炎': 'Large bowel_Colitis', '大肠_浆膜面毛糙': 'Large bowel_Serosal surface irregularity', '大肠_溃疡性结肠炎': 'Large bowel_Ulcerative colitis', '大肠_直肠癌': 'Large bowel_Rectal cancer', '大肠_积液积气': 'Large bowel_Gas and fluid accumulation', '大肠_结肠癌': 'Large bowel_Colon cancer', '大肠_肠壁毛糙': 'Large bowel_Wall irregularity', '大肠_肠壁水肿': 'Large bowel_Wall edema', '大肠_肠套叠': 'Large bowel_Intussusception', '大肠_肠憩室': 'Large bowel_Diverticulum', '大肠_肠梗阻': 'Large bowel_Obstruction', '大肠_肠穿孔': 'Large bowel_Perforation', '大肠_肠道扩张': 'Large bowel_Dilatation', '大肠_脂肪间隙模糊': 'Large bowel_Blurring of fat planes', '大肠_阑尾炎': 'Large bowel_Appendicitis', '大肠_阑尾粪石': 'Large bowel_Appendicolith', '小肠_克罗恩病': "Small bowel_Crohn's disease", '小肠_套叠': 'Small bowel_Intussusception', '小肠_扭转': 'Small bowel_Volvulus', '小肠_梗阻': 'Small bowel_Obstruction', '小肠_淋巴瘤': 'Small bowel_Lymphoma', '小肠_积气积液': 'Small bowel_Gas and fluid accumulation', '小肠_系膜指膜炎': 'Small bowel_Mesenteric panniculitis', '小肠_系膜淋巴结肿大': 'Small bowel_Mesenteric lymphadenopathy', '小肠_肠壁增厚': 'Small bowel_Wall thickening', '小肠_肠管扩张': 'Small bowel_Dilatation', '小肠_脂肪瘤': 'Small bowel_Lipoma', '小肠_间质瘤（胃肠间质瘤-gist）': 'Small bowel_Gastrointestinal stromal tumor', '小肠_（急慢性）小肠炎': 'Small bowel_Enteritis', '心脏_心包积液': 'Heart_Pericardial effusion', '心脏_心影（脏）增大': 'Heart_Cardiomegaly', '肋骨_转移瘤（乳腺癌 骨转移）': 'Rib_Metastasis', '肋骨_骨折': 'Rib_Fracture', '肋骨_骨质破坏': 'Rib_Bone destruction', '肝_低密度影': 'Liver_Hypoattenuating lesion', '肝_格林森鞘积液': 'Liver_Periportal edema', '肝_比例失调': 'Liver_Lobar volume disproportion', '肝_波浪状改变': 'Liver_Undulating contour', '肝_硬化': 'Liver_Cirrhosis', '肝_结节状强化': 'Liver_Nodular enhancement', '肝_肝内胆管扩张': 'Liver_Intrahepatic bile duct dilatation', '肝_肝内胆管结石': 'Liver_Hepatolithiasis', '肝_肝内钙化灶': 'Liver_Intrahepatic calcification', '肝_肝囊肿': 'Liver_Cyst', '肝_肝细胞癌': 'Liver_Hepatocellular carcinoma', '肝_肝胆管内高密度影': 'Liver_Hyperattenuating lesion in intrahepatic bile ducts', '肝_肝血管瘤': 'Liver_Hemangioma', '肝_胆管癌': 'Liver_Intrahepatic cholangiocarcinoma', '肝_脂肪肝': 'Liver_Steatotic liver disease', '肝_脓肿': 'Liver_Abscess', '肝_转移瘤': 'Liver_Metastasis', '肝_边缘不规则': 'Liver_Irregular margin', '肺_斑片影': 'Lung_Patchy opacity', '肺_气胸': 'Lung_Pneumothorax', '肺_结节': 'Lung_Nodule', '肺_肺占位': 'Lung_Mass', '肺_肺萎陷': 'Lung_Pulmonary collapse', '肺_胸腔积液': 'Lung_Pleural effusion', '肺_膨胀不全': 'Lung_Atelectasis', '肺_转移瘤': 'Lung_Metastasis', '肺_钙化灶': 'Lung_Calcification', '肺_高密度影': 'Lung_Hyperattenuating opacity', '肾_低密度影': 'Kidney_Hypoattenuating lesion', '肾_囊肿': 'Kidney_Cyst', '肾_多囊肾': 'Kidney_Polycystic kidney disease', '肾_实质变薄': 'Kidney_Parenchymal thinning', '肾_无强化囊性灶': 'Kidney_Nonenhancing cystic lesion', '肾_肾动脉瘤': 'Kidney_Renal artery aneurysm', '肾_肾盂扩张': 'Kidney_Renal pelvic dilatation', '肾_肾盂癌': 'Kidney_Renal pelvic cancer', '肾_肾盂积水': 'Kidney_Hydronephrosis', '肾_肾细胞癌（透明细胞癌）': 'Kidney_Renal cell carcinoma', '肾_肾萎缩': 'Kidney_Atrophy', '肾_肾血管平滑肌脂肪瘤': 'Kidney_Angiomyolipoma', '肾_肾（盂）结石': 'Kidney_Nephrolithiasis', '肾_高密度影': 'Kidney_Hyperattenuating lesion', '肾上腺_增生': 'Adrenal gland_Hyperplasia', '肾上腺_结节': 'Adrenal gland_Nodule', '肾上腺_脂肪瘤': 'Adrenal gland_Lipoma', '肾上腺_腺瘤': 'Adrenal gland_Adenoma', '肾上腺_转移瘤': 'Adrenal gland_Metastasis', '肾上腺_钙化': 'Adrenal gland_Calcification', '胃_壁水肿': 'Stomach_Wall edema', '胃_扩张': 'Stomach_Dilatation', '胃_胃底静脉曲张': 'Stomach_Gastric fundal varices', '胃_胃溃疡': 'Stomach_Ulcer', '胃_胃癌': 'Stomach_Gastric cancer', '胃_间质瘤（gist）': 'Stomach_Gastrointestinal stromal tumor (GIST)', '胆囊_结石': 'Gallbladder_Cholecystolithiasis', '胆囊_结节状致密影': 'Gallbladder_Nodular stone-like hyperattenuating lesion', '胆囊_胆囊增大': 'Gallbladder_Distention', '胆囊_胆囊炎': 'Gallbladder_Cholecystitis', '胆囊_胆囊癌': 'Gallbladder_Gallbladder cancer', '胆囊_胆囊腺肌症': 'Gallbladder_Adenomyomatosis', '胆囊_胆管壁增厚': 'Gallbladder_Extrahepatic bile duct wall thickening', '胆囊_胆管扩张': 'Gallbladder_Extrahepatic bile duct dilatation', '胆囊_胆管炎': 'Gallbladder_Cholangitis', '胆囊_胆管癌': 'Gallbladder_Cholangiocarcinoma', '胆囊_胆管积气': 'Gallbladder_Pneumobilia', '胆囊_胆管结石': 'Gallbladder_Extrahepatic bile duct stone', '胆囊_高密度影': 'Gallbladder_Hyperattenuating lesion', '胆囊_黄色肉芽肿': 'Gallbladder_Xanthogranuloma', '胰腺_低密度影': 'Pancreas_Low-density lesion', '胰腺_囊肿': 'Pancreas_Cyst', '胰腺_围脂肪间隙模糊': 'Pancreas_Blurring of peripancreatic fat planes', '胰腺_肿瘤或胰腺癌': 'Pancreas_Pancreatic cancer', '胰腺_胰周假性囊肿': 'Pancreas_Peripancreatic pseudocyst', '胰腺_胰管扩张': 'Pancreas_Pancreatic duct dilatation', '胰腺_胰管结石': 'Pancreas_Pancreatic duct calculus', '胰腺_胰腺炎': 'Pancreas_Pancreatitis', '胰腺_胰腺饱满': 'Pancreas_Enlargement', '胰腺_萎缩': 'Pancreas_Atrophy', '脾_低密度灶': 'Spleen_Hypoattenuating lesion', '脾_副脾': 'Spleen_Accessory spleen', '脾_囊肿': 'Spleen_Cyst', '脾_梗死': 'Spleen_Infarction', '脾_片状低密度区': 'Spleen_Patchy hypoattenuating lesion', '脾_脾大': 'Spleen_Splenomegaly', '脾_脾脏淋巴瘤': 'Spleen_Lymphoma', '脾_钙化': 'Spleen_Calcification', '膀胱_憩室': 'Bladder_Diverticulum', '膀胱_结石': 'Bladder_Stone', '膀胱_膀胱壁毛糙': 'Bladder_Wall irregularity', '膀胱_膀胱炎': 'Bladder_Cystitis', '膀胱_膀胱癌': 'Bladder_Bladder cancer', '膀胱_软组织密度影': 'Bladder_Soft-tissue attenuation lesion', '门静脉_增宽': 'Portal vein_Dilatation', '门静脉_栓塞': 'Portal vein_Thrombosis', '门静脉_高压': 'Portal vein_Hypertension', '食管_增粗迂曲血管影': 'Esophagus_Dilated and tortuous tubular opacities', '食管_管壁增厚': 'Esophagus_Wall thickening', '食管_裂孔疝': 'Esophagus_Hiatal hernia', '食管_静脉扩张迂曲': 'Esophagus_Dilated and tortuous veins', '食管_静脉曲张': 'Esophagus_Varices', '骶骨_骨炎': 'Sacrum_Osteitis'}
+        self.test_organs = list(set([item.split('_')[0] for item in self.test_items]))
+    
+    def __len__(self):
+        return len(self.img_paths)
+
+    def __getitem__(self, index):
+        # load image
+        image_path = self.img_paths[index]
+        data = {"image": image_path}
+        res = transforms.LoadImaged(keys=["image"], image_only=False, ensure_channel_first=True)(data)
+        image = res["image"]
+        
+        affine = res["image_meta_dict"]["affine"]
+        spacing = (
+            abs(affine[0, 0].item()),
+            abs(affine[1, 1].item()),
+            abs(affine[2, 2].item())
+        )
+        _, h, w, d = image.shape
+        orig_shape_hwd = (h, w, d)
+
+        ref_spacing = (1.0, 1.0, 5.0)
+        scale = [spacing[i] / ref_spacing[i] for i in range(3)]
+        target_size = [int(h * scale[1]), int(w * scale[0]), int(d * scale[2])]  # [H', W', D']
+
+        trans = transforms.Compose(
+            [
+                transforms.Resized(spatial_size=target_size, keys=["image"], mode="trilinear"),
+                transforms.Transposed(keys=["image"], indices=(0, 3, 2, 1)),
+            ]
+        )
+        resized_data = trans(res)
+
+        img_resized = resized_data["image"]   # [C, D', W', H']
+        image = img_resized
+        image[image > 400] = 400
+        image[image < -300] = -300
+        image = (image - image.min()) / (image.max() - image.min() + 1e-8)
+        img = image
+
+        # crop non-zero region in image
+        roi_coords = np.nonzero(img[0].cpu().numpy())
+        min_dhw = torch.from_numpy(np.min(roi_coords, axis=1))
+        max_dhw = torch.from_numpy(np.max(roi_coords, axis=1))
+
+        extend_d = 5
+        extend_hw = 20
+
+        min_dhw = torch.max(
+            min_dhw - torch.tensor([extend_d, extend_hw, extend_hw]),
+            torch.tensor([0, 0, 0]),
+        )
+        max_dhw = torch.min(
+            max_dhw + torch.tensor([extend_d, extend_hw, extend_hw]),
+            torch.tensor([img.shape[1], img.shape[2], img.shape[3]]),
+        )
+
+        cropped_image = img[
+            :,
+            min_dhw[0]: max_dhw[0],
+            min_dhw[1]: max_dhw[1],
+            min_dhw[2]: max_dhw[2]
+        ]
+        crop_shape_dhw = tuple(cropped_image.shape[1:])
+
+        # pad data to [96, 256, 384] if smaller
+        data["image"] = cropped_image
+        data_pad = self.pad_func(data)
+        data = data_pad
+
+        file_name = image_path.split('/')[-1]
+        patient_id = file_name.split('_')[0]
+        test_organ_names = self.test_organs
+        
+        meta_info = {
+            'file_name': file_name,
+            'img_path': image_path,
+            'patient_id': patient_id,
+            'test_organ_names': test_organ_names,
+            'letter': 'None',
+        }
+        return data['image'].as_tensor(), self.test_items, meta_info
+
+
+class RADAR(nn.Module):
+    def __init__(
+        self,
+        image_encoder,
+        text_encoder,
+        text_decoder=None,
+        queue_size=1234,
+        alpha=0.4,
+        embed_dim=256,
+        momentum=0.995,
+        tie_enc_dec_weights=True,
+        max_txt_len=175,
+    ):
+        super().__init__()
+
+        self.tokenizer = BertTokenizer.from_pretrained(os.path.join(configs_root, "bert-base-chinese"))
+
+        text_encoder.resize_token_embeddings(len(self.tokenizer))
+        self.visual_encoder = image_encoder
+        self.text_encoder = text_encoder
+
+        text_width = text_encoder.config.hidden_size
+        vision_width = 256
+
+        self.text_proj = nn.Linear(text_width, embed_dim)
+
+        self.queue_size = queue_size
+        self.momentum = momentum
+        self.temp = nn.Parameter(0.07 * torch.ones([]))
+
+        self.alpha = alpha
+        self.max_txt_len = max_txt_len
+
+        self.organs = [
+            '肾上腺', '主动脉', '竖脊肌', '脑', '锁骨', '大肠', '十二指肠', '食管', '面部', '股骨', 
+            '胆囊', "臀肌", '心脏', '髋关节', '肱骨', '髂动脉', '髂静脉', '髂腰肌', '下腔静脉', '肾', 
+            '肝', '肺', '胰腺', '门静脉', '肺动脉', '肋骨', '骶骨', '肩胛骨', '小肠', '脾', 
+            '胃', '气管', '膀胱', '颈椎', '腰椎', '胸椎'
+        ]
+        # organ_dict = {
+        #     "肾上腺": "adrenal gland", "主动脉": "aorta", "竖脊肌": "erector spinae muscle", "脑": "brain", "锁骨": "clavicle", "大肠": "large bowel", "十二指肠": "duodenum", 
+        #     "食管": "esophagus", "面部": "face", "股骨": "femur", "胆囊": "gallbladder", "臀肌": "gluteus muscle", "心脏": "heart", "髋关节": "hip joint", "肱骨": "humerus", 
+        #     "髂动脉": "iliac artery", "髂静脉": "iliac vena", "髂腰肌": "iliopsoas muscle", "下腔静脉": "inferior vena cava", "肾": "kidney", "肝": "liver", "肺": "lung", 
+        #     "胰腺": "pancreas", "门静脉": "portal vein", "肺动脉": "pulmonary artery", "肋骨": "rib", "骶骨": "sacrum", "肩胛骨": "scapula", "小肠": "small bowel", "脾": "spleen", 
+        #     "胃": "stomach", "气管": "trachea", "膀胱": "bladder", "颈椎": "cervical vertebrae", "腰椎": "lumbar vertebrae", "胸椎": "thoracic vertebrae"
+        # }
+        
+        self.attention = nn.MultiheadAttention(
+            embed_dim=vision_width,
+            num_heads=4,
+            dropout=0.1,
+            batch_first=True
+        )
+
+        self.vision_projs = nn.ModuleList([nn.Linear(vision_width, embed_dim) for _ in range(len(self.organs))])
+        self.query_tokens = nn.Parameter(torch.zeros(len(self.organs), vision_width))
+
+    @torch.inference_mode()
+    def forward_test_win(
+        self, 
+        images, 
+        masks, 
+        organ_logits,
+        test_organs,
+        text_feat_dict,
+        organ_feat_dict,
+        whole_organ_sizes,
+        skip_organ=None
+    ):
+        seg_probs, seg, image_embeds1, image_embeds2, image_embeds3, organ_token_flags1, organ_token_flags2, organ_token_flags3 = self.visual_encoder(images, None)
+
+        margin = 2
+        masks = seg
+        
+        for i, (embed1, embed2, embed3, mask) in enumerate(zip(image_embeds1, image_embeds2, image_embeds3, masks)):
+            boundaries = []
+            for d in range(mask.dim()):
+                start_slice = [slice(None)] * mask.dim()
+                end_slice = [slice(None)] * mask.dim()
+                
+                start_slice[d] = slice(None, margin)
+                end_slice[d] = slice(-margin, None)
+                
+                boundaries.append(mask[tuple(start_slice)][mask[tuple(start_slice)] > 0])
+                boundaries.append(mask[tuple(end_slice)][mask[tuple(end_slice)] > 0])
+            boundaries = torch.cat(boundaries)
+            
+            boundary_values = boundaries[boundaries > 0].flatten()
+            boundary_organs = torch.unique(boundary_values)
+
+            if skip_organ is not None:
+                boundary_organs = boundary_organs[boundary_organs != skip_organ + 1]
+            
+            organ_ids, organ_counts = torch.unique(mask, return_counts=True)
+            organ_ids = organ_ids.long()
+            organ_counts = organ_counts[organ_ids != 0]
+            organ_ids = organ_ids[organ_ids != 0]
+
+            # organs not touch boundary
+            intact_organ_ids = [organ_id for organ_id, organ_count in zip(organ_ids, organ_counts) if organ_id not in boundary_organs]
+            intact_organ_ids = torch.tensor(intact_organ_ids, device=masks.device).long()
+            intact_organ_ids = intact_organ_ids - 1
+            
+            if not len(intact_organ_ids):
+                continue
+
+            organ_sizes = dict(zip([self.organs[organ_id] for organ_id in intact_organ_ids], [organ_counts[organ_ids == organ_id + 1].item() for organ_id in intact_organ_ids]))
+
+            for organ_id in intact_organ_ids:
+                organ_name = self.organs[organ_id.item()]
+                if organ_name not in test_organs:
+                    continue
+                    
+                if organ_name in organ_feat_dict:
+                    continue
+
+                tokens1 = organ_token_flags1[i, organ_id, :]
+                tokens2 = organ_token_flags2[i, organ_id, :]
+                tokens3 = organ_token_flags3[i, organ_id, :]
+
+                query = self.query_tokens[organ_id].unsqueeze(0).unsqueeze(0)
+                key1 = embed1[tokens1].unsqueeze(0)
+                key2 = embed2[tokens2].unsqueeze(0)
+                key3 = embed3[tokens3].unsqueeze(0)
+
+                key = value = torch.cat([key1, key2, key3], dim=1)
+                
+                updated_query_token, _ = self.attention(query, key, value)
+                updated_query_token = updated_query_token.squeeze(0)
+
+                image_feat = F.normalize(self.vision_projs[organ_id](updated_query_token), dim=-1)
+                
+                organ_feat_dict[organ_name] = image_feat.cpu().tolist()
+
+                for item in organ_logits.keys():
+                    if isinstance(item, str):
+                        item_organ_name = item.split('_')[0]
+                    else:
+                        item_organ_name = item[0]
+                    if item_organ_name != organ_name:
+                        continue
+
+                    text_feat = text_feat_dict[item]
+
+                    logits = image_feat @ text_feat.t() / self.temp
+                    probs = logits.softmax(-1)
+                    organ_logits[item].append(probs.cpu().tolist())
+    
+        return organ_logits, seg_probs
+
+@torch.inference_mode()
+def evaluate(pad_func, model, img_dir, save_dir, save_tag):
+
+    datafolder = DataFolder(img_dir)
+    dataloader = DataLoader(
+        datafolder,
+        batch_size=1,
+        shuffle=False,
+        num_workers=0,  # Space patch: single-volume inference; avoids forking 12 torch workers in the CPU container.
+        drop_last=False,
+        collate_fn=collate_fn
+    )
+
+    sw_batch_size = 1
+    overlap = 0.25
+    roi_size = (96, 256, 384)
+
+    miss_num = 0
+    results = []
+    organ_status = {}
+
+    # load pos/neg ensembled prompt embeddings
+    # Space patch: resolve via MODEL_ROOT (absolute in the Space) instead of cwd-relative '../ckpt'.
+    text_feat_dict = torch.load(os.path.join(model_root, 'infer_text_embedding_radar.pt'))
+    organ_feat_dict = {}
+    save_path = os.path.join(save_dir, f'RADAR_infer_results_{save_tag}.csv')
+    os.makedirs(save_dir, exist_ok=True)
+    
+    for i, (image, test_items, meta_info) in enumerate(tqdm(dataloader, desc='Infer')):
+        torch.cuda.empty_cache()
+        skip_case = False
+        for tmp_s in image.shape[1:]:
+            if tmp_s > 1000:
+                skip_case = True
+                break
+        if skip_case:
+            continue
+        
+        fid = meta_info['file_name']
+        organ_feat_dict[fid] = {}
+
+        image = image[None].cuda()
+
+        test_organs = meta_info['test_organ_names']
+
+        image_size = list(image.shape[2:])
+        num_spatial_dims = len(image.shape) - 2
+
+        scan_interval = _get_scan_interval(
+            image_size, roi_size, num_spatial_dims, overlap
+        )
+        slices = dense_patch_slices(image_size, roi_size, scan_interval)
+        num_win = len(slices)
+        organ_logits = dict(zip(test_items, [[] for _ in test_items]))
+        # organ_logits.pop('胆囊_术后胆囊缺失')  # surgically_absent_gallbladder
+
+        # get full mask
+        full_mask = torch.zeros((1, 37) + tuple(image_size)).cuda()
+        count_map = torch.zeros_like(full_mask).cuda()
+
+        for slice_g in range(0, num_win, sw_batch_size):
+            slice_range = range(slice_g, min(slice_g + sw_batch_size, num_win))
+            unravel_slice = [
+                [slice(int(idx / num_win), int(idx / num_win) + 1), slice(None)] + list(slices[idx % num_win])
+                for idx in slice_range
+            ]
+            
+            window_patches = torch.cat([image[win_slice] for win_slice in unravel_slice]).cuda()
+
+            organ_logits, pred_window_seg_prob = model.forward_test_win(
+                window_patches, 
+                None,
+                organ_logits,
+                test_organs,
+                text_feat_dict,
+                organ_feat_dict[fid],
+                None
+            )
+
+            # interpolate
+            interpolated_seg_prob = F.interpolate(pred_window_seg_prob, size=window_patches.shape[2:], mode='trilinear')
+            
+            for ii, slice_idx in enumerate(slice_range):
+                full_slice = unravel_slice[ii]
+                full_mask[full_slice] += interpolated_seg_prob[ii]
+                count_map[full_slice] += 1
+        
+        # Avoid division by zero by ensuring count_map is at least 1 everywhere
+        count_map = torch.clamp(count_map, min=1)
+        stitched_mask = full_mask / count_map  # argmax
+        stitched_mask = stitched_mask.argmax(1).unsqueeze(0)
+    
+        margin = 2
+        boundaries = []
+        squeeze_stitched_mask = stitched_mask.squeeze(0).squeeze(0)
+        for d in range(squeeze_stitched_mask.dim()):
+            start_slice = [slice(None)] * squeeze_stitched_mask.dim()
+            end_slice = [slice(None)] * squeeze_stitched_mask.dim()
+            
+            start_slice[d] = slice(None, margin)
+            end_slice[d] = slice(-margin, None)
+            
+            boundaries.append(squeeze_stitched_mask[tuple(start_slice)][squeeze_stitched_mask[tuple(start_slice)] > 0])
+            boundaries.append(squeeze_stitched_mask[tuple(end_slice)][squeeze_stitched_mask[tuple(end_slice)] > 0])
+        boundaries = torch.cat(boundaries)
+        
+        boundary_values = boundaries[boundaries > 0].flatten()
+        boundary_organs = torch.unique(boundary_values)
+        
+        organ_ids, organ_counts = torch.unique(squeeze_stitched_mask, return_counts=True)
+        organ_ids = organ_ids.long()
+        organ_counts = organ_counts[organ_ids != 0]
+        organ_ids = organ_ids[organ_ids != 0]
+
+        # organs not touch boundary
+        intact_organ_ids = [organ_id for organ_id, organ_count in zip(organ_ids, organ_counts) if organ_id not in boundary_organs]
+        intact_organ_ids = torch.tensor(intact_organ_ids, device=squeeze_stitched_mask.device).long()
+        intact_organ_ids = intact_organ_ids - 1
+
+        # for melrin data, we just infer all organs
+        # organ_logits = {k:v for k,v in organ_logits.items() if datafolder.organs.index(k[0]) in intact_organ_ids}
+
+        for k, v in organ_logits.items():
+            if not len(v):
+                organ_name = k.split('_')[0]
+                organ_id = datafolder.organs.index(organ_name)
+
+                window_patch, window_mask = center_crop(
+                    image,
+                    torch.eq(stitched_mask, organ_id + 1),
+                    crop_size=roi_size
+                )
+                window_mask = window_mask.float()
+                window_mask[window_mask == 1] = organ_id + 1
+
+                pad_data = pad_func({'image': window_patch[0], 'label': window_mask[0]})
+                window_patch, window_mask = pad_data['image'], pad_data['label']
+
+                organ_logits, _ = model.forward_test_win(
+                    window_patch[None], 
+                    None,
+                    organ_logits,
+                    test_organs,
+                    text_feat_dict,
+                    organ_feat_dict[fid],
+                    None,
+                    skip_organ=organ_id
+                )
+  
+        res = [meta_info['file_name']] + [''] * len(datafolder.test_items)
+        organ_logits = {item: probs for item, probs in organ_logits.items() if len(probs) > 0}
+        
+        for item, probs in organ_logits.items():
+            res[datafolder.test_items.index(item) + 1] = np.concatenate(probs).mean(0)[1]  # get average of one organ in multi-widows
+        results.append(res)
+    
+    if dist.is_initialized():
+        results = np.concatenate(all_gather(results), axis=0)
+    else:
+        results = results
+    
+    pd.DataFrame(
+        results,
+        columns=['file_name'] + [f'{k} ({datafolder.english_mapping[k]})' for k in datafolder.test_items]
+    ).to_csv(save_path, index=False, encoding='utf-8-sig')
+
+def initialize():
+    """
+    Returns: transforms.DivisiblePadd, RADAR
+    """
+    print('\n--> Start initializing...')
+    pad_func = transforms.DivisiblePadd(
+        keys=["image", "label"],
+        k=32,
+        mode='constant',
+        constant_values=0,
+        method="end"
+    )
+
+    vision_encoder = VisionBranch()
+    text_encoder = XBertEncoder.from_config({}, from_pretrained=True)
+
+    model = RADAR(
+        image_encoder=vision_encoder,
+        text_encoder=text_encoder,
+    )
+
+    ckpt_path = os.path.join(model_root, "checkpoint_radar_pretrain.pth")
+    print('--> ckpt_path: ', ckpt_path)
+    ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=False)
+
+    msg = model.load_state_dict(ckpt['model'], strict=False)
+
+    model.eval()
+    model.cuda()
+
+    print('\n--> Initialize done')
+
+    return pad_func, model
+
+def inference(initialize_returns, img_dir, save_dir, save_tag):
+    """
+    Args:
+        initialize_returns: pad_func, model
+        img_dir: see argparse
+        save_dir: see argparse
+    """
+    print('\n--> Start inference.')
+    pad_func, model = initialize_returns
+    evaluate(pad_func, model, img_dir, save_dir, save_tag)
+    csv_file = os.path.join(save_dir, f'RADAR_infer_results_{save_tag}.csv')
+    print(f'evaluate done, save result_csv to {csv_file}.')
+    
+    # TODO: compute metrics
+    
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--img_dir', type=str, default='../data/demo_cases', help='The path to inference image folder.')
+    parser.add_argument('--save_dir', type=str, default='../results', help='The path to save folder.')
+    parser.add_argument('--save_tag', type=str, default='demo', help='Save tag.')
+    
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    initialize_returns = initialize()
+    
+    # infer
+    inference(initialize_returns, args.img_dir, args.save_dir, args.save_tag)
+
+
+if __name__ == '__main__':
+    main()
+
+
+

README.md

@@ -0,0 +1,21 @@
+---
+title: RADAR Abdominal CT
+emoji: 🩻
+colorFrom: blue
+colorTo: green
+sdk: gradio
+app_file: app.py
+pinned: false
+license: cc-by-nc-sa-4.0
+python_version: "3.12"
+preload_from_hub:
+  - radar-generalist/RADAR checkpoint_radar_pretrain.pth,bert-base-chinese/config.json,bert-base-chinese/config_decoder.json,bert-base-chinese/pytorch_model.bin,bert-base-chinese/tokenizer.json,bert-base-chinese/tokenizer_config.json,bert-base-chinese/vocab.txt
+---
+
+RADAR abdominal-CT findings demo: upload a contrast-enhanced abdominal CT volume (`.nii` / `.nii.gz`), get per-finding scores plus a ranked report.
+
+- Paper/model: [radar-generalist/RADAR](https://huggingface.co/radar-generalist/RADAR) (CC BY-NC-SA 4.0, non-commercial research use only)
+- Source: vendored inference code from the RADAR release (Zenodo `damo-radar.zip`, record `21271172`); weights load at runtime from the public HuggingFace repo, no token needed
+- Runtime: ZeroGPU (`large`), single `@spaces.GPU(duration=180)` call per volume; preprocessing (1x1x5mm resample, 96x256x384 pad/crop) is the upstream MONAI chain, reused verbatim
+
+Research assistance tool, not a medical diagnosis.

app.py

@@ -0,0 +1,133 @@
+"""RADAR abdominal-CT ZeroGPU Space.
+
+Reuse map (all inference logic is vendored, not reimplemented):
+- `RADAR_inference/inference_demo.py::initialize` builds the RADAR model and
+  loads `ckpt/checkpoint_radar_pretrain.pth` (strict=False), then `.cuda()`.
+- `RADAR_inference/inference_demo.py::evaluate` runs the full single-volume
+  pipeline: MONAI resample to 1x1x5mm -> HU clip [-300,400] -> min-max norm ->
+  non-zero ROI crop -> pad (96,256,384) -> sliding-window forward with
+  `inference_demo.RADAR.forward_test_win` -> per-organ center-crop second
+  pass -> CSV of mean positive-class scores for 146 organ_finding pairs.
+- `RADAR_inference/inference_demo.py::DataFolder` owns every preprocessing
+  transform. This file adds no transforms and no report prose: it bridges
+  Gradio upload -> temp dir -> evaluate -> (Label, Dataframe of raw scores).
+
+Space layout notes:
+- `MODEL_ROOT`/`CONFIGS_ROOT` must be absolute before importing the vendored
+  module (it reads them at import time). Weights arrive preloaded at build
+  time (`preload_from_hub`, same HF cache `snapshot_download` reads) and are
+  symlinked into `ckpt/`; prompt embeddings (`infer_text_embedding_radar.pt`,
+  340KB) ship in git because they are absent from the HuggingFace repo.
+"""
+
+import os
+import re
+import shutil
+import tempfile
+
+ROOT = os.path.dirname(os.path.abspath(__file__))
+CKPT_DIR = os.path.join(ROOT, "ckpt")
+os.environ.setdefault("MODEL_ROOT", CKPT_DIR)
+os.environ.setdefault("CONFIGS_ROOT", CKPT_DIR)
+
+import sys
+sys.path.insert(0, os.path.join(ROOT, "RADAR_inference"))  # vendored absolute imports (dynamic_network_architectures) resolve from here
+import pandas as pd  # noqa: E402
+import spaces  # noqa: E402
+import gradio as gr  # noqa: E402
+from huggingface_hub import snapshot_download  # noqa: E402
+from inference_demo import initialize, evaluate  # noqa: E402
+
+REPO_ID = "radar-generalist/RADAR"
+CHECKPOINT_NAME = "checkpoint_radar_pretrain.pth"
+
+
+def _ensure_ckpts() -> None:
+    snap = snapshot_download(
+        REPO_ID,
+        allow_patterns=[CHECKPOINT_NAME, "bert-base-chinese/*"],
+    )
+    targets = [CHECKPOINT_NAME, "bert-base-chinese"]
+    for name in targets:
+        src = os.path.join(snap, name)
+        dst = os.path.join(CKPT_DIR, name)
+        if not os.path.exists(src):
+            raise RuntimeError(f"{name} absent from {REPO_ID} snapshot {snap}")
+        if os.path.lexists(dst):
+            continue
+        os.symlink(src, dst)
+    ckpt_path = os.path.join(CKPT_DIR, CHECKPOINT_NAME)
+    if not os.path.exists(ckpt_path):
+        raise RuntimeError(
+            f"{CHECKPOINT_NAME} missing after download; "
+            "check Space logs/network and re-run."
+        )
+
+
+_ensure_ckpts()
+pad_func, model = initialize()  # module level; .cuda() here is intentional (ZeroGPU pattern)
+
+
+def _english_name(column: str) -> str:
+    m = re.search(r"\((.+)\)\s*$", column)
+    return m.group(1) if m else column
+
+
+@spaces.GPU(duration=180)
+def diagnose(ct_file):
+    """Run RADAR on one uploaded contrast-enhanced abdominal CT volume."""
+    path = ct_file if isinstance(ct_file, str) else ct_file.name
+    if not path.endswith((".nii", ".nii.gz")):
+        raise gr.Error("upload a .nii or .nii.gz CT volume")
+    try:
+        import nibabel as nib
+
+        nib.load(path)
+    except Exception:
+        raise gr.Error("upload a .nii or .nii.gz CT volume")
+
+    tmpdir = tempfile.mkdtemp(prefix="radar_case_")
+    outdir = tempfile.mkdtemp(prefix="radar_out_")
+    try:
+        fname = os.path.basename(path)
+        if not fname.endswith((".nii", ".nii.gz")):
+            fname += ".nii.gz"
+        shutil.copy(path, os.path.join(tmpdir, fname))
+        try:
+            evaluate(pad_func, model, tmpdir, outdir, "space")
+        except (OSError, ValueError, RuntimeError) as exc:
+            raise gr.Error(f"could not process this volume: {exc}")
+        csv_path = os.path.join(outdir, "RADAR_infer_results_space.csv")
+        df = pd.read_csv(csv_path, encoding="utf-8-sig")
+        if df.empty:
+            raise gr.Error("model skipped this volume (check dimensions/spacing)")
+        row = df.iloc[0]
+        scores = {
+            _english_name(col): float(row[col])
+            for col in df.columns[1:]
+            if str(row[col]).strip() != ""
+        }
+        ranked = sorted(scores.items(), key=lambda kv: kv[1], reverse=True)
+        table = pd.DataFrame(ranked, columns=["Finding", "Score"])
+        return scores, table
+    finally:
+        shutil.rmtree(tmpdir, ignore_errors=True)
+        shutil.rmtree(outdir, ignore_errors=True)
+
+
+demo = gr.Interface(
+    fn=diagnose,
+    inputs=gr.File(
+        label="Abdominal CT (.nii.gz)", file_types=[".nii", ".nii.gz"]
+    ),
+    outputs=[gr.Label(label="Top findings", num_top_classes=10), gr.Dataframe(label="All finding scores")],
+    title="RADAR Abdominal CT",
+    description=(
+        "Expert-level generalist AI for contrast-enhanced abdominal CT "
+        "(18 structures, 146 findings). Non-commercial research demo "
+        "(CC BY-NC-SA 4.0); assistance tool, not a diagnosis."
+    ),
+)
+
+if __name__ == "__main__":
+    demo.launch()

ckpt/infer_text_embedding_radar.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ac960b0300c2ba10526e3bf500cd0c879fe5d9d15b9c49c17c34f0fc40e6cbe
+size 346348

requirements.txt

@@ -0,0 +1,45 @@
+contexttimer
+decord
+diffusers<=0.16.0
+einops>=0.4.1
+fairscale==0.4.4
+ftfy
+iopath
+ipython
+omegaconf
+# opencv 4.5.5.64 ships no cp312 files: keep upstream pin below 3.12 only.
+opencv-python-headless==4.5.5.64; python_version < '3.12'
+opencv-python-headless; python_version >= '3.12'
+opendatasets
+packaging
+pandas
+plotly
+pre-commit
+pycocoevalcap
+pycocotools
+python-magic
+scikit-image
+sentencepiece
+spacy
+streamlit
+timm==0.4.12
+torch>=1.10.0
+torchvision
+tqdm
+# transformers 4.25 pulls tokenizers 0.13.3 (no cp312 wheel, needs Rust): keep it below 3.12 only.
+transformers==4.25; python_version < '3.12'
+transformers==4.48.3; python_version >= '3.12'
+webdataset
+wheel
+h5py
+monai
+batchgenerators
+SimpleITK
+nibabel
+matplotlib
+nltk
+gradio>=4
+huggingface_hub
+nibabel
+monai
+torch>=2.8,<2.14