Method __init__

lib/networks.py:633–676

(self, n_class=1, img_size_s1=(256,256), img_size_s2=(224,224), decoder_aggregation='additive', interpolation='bilinear')

                        
class MERIT_Parallel_Small(nn.Module):
    def __init__(self, n_class=1, img_size_s1=(256,256), img_size_s2=(224,224), decoder_aggregation='additive', interpolation='bilinear'):
        super(MERIT_Parallel_Small, self).__init__()
        
        self.interpolation = interpolation
        self.img_size_s1 = img_size_s1
        self.img_size_s2 = img_size_s2

        # conv block to convert single channel to 3 channels
        self.conv = nn.Sequential(
            nn.Conv2d(1, 3, kernel_size=1),
            nn.BatchNorm2d(3),
            nn.ReLU(inplace=True)
        )
        
        # backbone network initialization with pretrained weight
        #self.backbone = maxvit_tiny_rw_224_4out()  # [64, 128, 320, 512]
        self.backbone1 = maxxvit_rmlp_small_rw_256_4out()  # [64, 128, 320, 512]
        self.backbone2 = maxvit_rmlp_small_rw_224_4out()  # [64, 128, 320, 512]
        
        print('Loading:', './pretrained_pth/maxvit/maxxvit_rmlp_small_rw_256_sw-37e217ff.pth')
        state_dict1 = torch.load('./pretrained_pth/maxvit/maxxvit_rmlp_small_rw_256_sw-37e217ff.pth')        
        self.backbone1.load_state_dict(state_dict1, strict=False)
        
        print('Loading:', './pretrained_pth/maxvit/maxvit_rmlp_small_rw_224_sw-6ef0ae4f.pth')       
        state_dict2 = torch.load('./pretrained_pth/maxvit/maxvit_rmlp_small_rw_224_sw-6ef0ae4f.pth')        
        self.backbone2.load_state_dict(state_dict2, strict=False)
        
        print('Pretrain weights loaded.')
        
        #channels=[512, 256, 128, 64]
        channels = [768, 384, 192, 96]
        # decoder initialization
        if(decoder_aggregation=='additive'):
            self.decoder = CASCADE_Add(channels=channels)
        elif(self.decoder_aggregation=='concatenation'):
            self.decoder = CASCADE_Cat(channels=self.channels)
        else:
            sys.exit("'"+self.decoder_aggregation+"' is not a valid decoder aggregation! Currently supported aggregations are 'additive' and 'concatenation'.")
        
        # Prediction heads initialization
        self.out_head1 = nn.Conv2d(channels[0], n_class, 1)
        self.out_head2 = nn.Conv2d(channels[1], n_class, 1)
        self.out_head3 = nn.Conv2d(channels[2], n_class, 1)
        self.out_head4 = nn.Conv2d(channels[3], n_class, 1)

    def forward(self, x, im1_size=(256,256), im2_size=(224,224)):