add model/dataset registry

Author: zh320

README.md

@@ -7,7 +7,7 @@ PyTorch implementation of realtime semantic segmentation models, support multi-g
 # Requirements
 
 torch == 1.8.1  
-segmentation-models-pytorch  
+segmentation-models-pytorch (optional)  
 torchmetrics  
 albumentations  
 loguru  

configs/my_config.py

@@ -3,7 +3,7 @@
 
 class MyConfig(BaseConfig):
     def __init__(self,):
-        super(MyConfig, self).__init__()
+        super().__init__()
         # Dataset
         self.dataset = 'cityscapes'
         self.data_root = '/path/to/your/dataset'

configs/optuna_config.py

@@ -7,7 +7,7 @@
 
 class OptunaConfig(BaseConfig):
     def __init__(self,):
-        super(OptunaConfig, self).__init__()
+        super().__init__()
         # Dataset
         self.dataset = 'cityscapes'
         self.data_root = '/path/to/your/dataset'

configs/parser.py

@@ -1,5 +1,8 @@
 import argparse
 
+from datasets import list_available_datasets
+from models import list_available_models
+
 
 def load_parser(config):
     args = get_parser()
@@ -16,7 +19,8 @@ def load_parser(config):
 def get_parser():
     parser = argparse.ArgumentParser()
     # Dataset
-    parser.add_argument('--dataset', type=str, default=None, choices=['cityscapes'],
+    dataset_list = list_available_datasets()
+    parser.add_argument('--dataset', type=str, default=None, choices=dataset_list,
         help='choose which dataset you want to use')
     parser.add_argument('--dataroot', type=str, default=None, 
         help='path to your dataset')
@@ -26,14 +30,8 @@ def get_parser():
         help='ignore index used for cross_entropy/ohem loss')
 
     # Model
-    parser.add_argument('--model', type=str, default=None, 
-        choices=['adscnet', 'aglnet', 'bisenetv1', 'bisenetv2', 'canet', 'cfpnet',
-                 'cgnet', 'contextnet', 'dabnet', 'ddrnet', 'dfanet', 'edanet',
-                 'enet', 'erfnet', 'esnet', 'espnet', 'espnetv2', 'fanet', 'farseenet',
-                 'fastscnn', 'fddwnet', 'fpenet', 'fssnet', 'icnet', 'lednet',
-                 'linknet', 'lite_hrnet', 'liteseg', 'mininet', 'mininetv2', 'ppliteseg', 
-                 'regseg', 'segnet', 'shelfnet', 'sqnet', 'stdc', 'swiftnet',
-                 'smp'],
+    model_list = list_available_models()
+    parser.add_argument('--model', type=str, default=None, choices=model_list,
         help='choose which model you want to use')
     parser.add_argument('--encoder', type=str, default=None, 
         help='choose which encoder of SMP model you want to use (please refer to SMP repo)')
@@ -179,4 +177,4 @@ def get_parser():
         help='temperature used for KL divergence loss')
 
     args = parser.parse_args()
-    return args
+    return args

core/__init__.py

@@ -1,3 +1,2 @@
-from .base_trainer import BaseTrainer
 from .seg_trainer import SegTrainer
 from .loss import get_loss_fn, kd_loss_fn, get_detail_loss_fn

core/base_trainer.py

@@ -12,7 +12,7 @@
 
 class BaseTrainer:
     def __init__(self, config):
-        super(BaseTrainer, self).__init__()
+        super().__init__()
         # DDP parameters, DO NOT CHANGE
         self.rank = int(os.getenv('RANK', -1))
         self.local_rank = int(os.getenv('LOCAL_RANK', -1))

core/loss.py

@@ -5,7 +5,7 @@
 
 class OhemCELoss(nn.Module):
     def __init__(self, thresh, ignore_index=255):
-        super(OhemCELoss, self).__init__()
+        super().__init__()
         self.thresh = -torch.log(torch.tensor(thresh, requires_grad=False, dtype=torch.float)).cuda()
         self.ignore_index = ignore_index
         self.criteria = nn.CrossEntropyLoss(ignore_index=ignore_index, reduction='none')
@@ -22,7 +22,7 @@ def forward(self, logits, labels):
 
 class DiceLoss(nn.Module):
     def __init__(self, smooth=1):
-        super(DiceLoss, self).__init__()
+        super().__init__()
         self.smooth = smooth
 
     def forward(self, logits, labels):
@@ -39,7 +39,7 @@ class DetailLoss(nn.Module):
     '''Implement detail loss used in paper
        `Rethinking BiSeNet For Real-time Semantic Segmentation`'''
     def __init__(self, dice_loss_coef=1., bce_loss_coef=1., smooth=1):
-        super(DetailLoss, self).__init__()
+        super().__init__()
         self.dice_loss_coef = dice_loss_coef
         self.bce_loss_coef = bce_loss_coef
         self.dice_loss_fn = DiceLoss(smooth)

datasets/__init__.py

@@ -1,7 +1,7 @@
 from torch.utils.data import DataLoader
-from .cityscapes import Cityscapes
 
-dataset_hub = {'cityscapes':Cityscapes,}
+from .cityscapes import Cityscapes
+from .dataset_registry import dataset_hub
 
 
 def get_dataset(config):
@@ -58,4 +58,10 @@ def get_test_loader(config):
         test_loader = DataLoader(dataset, batch_size=config.test_bs, 
                                     shuffle=False, num_workers=config.num_workers)
 
-    return test_loader
+    return test_loader
+
+
+def list_available_datasets():
+    dataset_list = list(dataset_hub.keys())
+
+    return dataset_list

datasets/cityscapes.py

@@ -5,9 +5,12 @@
 from torch.utils.data import Dataset
 import albumentations as AT
 from albumentations.pytorch import ToTensorV2
+
 from utils import transforms
+from .dataset_registry import register_dataset
 
 
+@register_dataset
 class Cityscapes(Dataset):
     # Codes are based on https://github.com/mcordts/cityscapesScripts
 

datasets/dataset_registry.py

@@ -0,0 +1,6 @@
+dataset_hub = {}
+
+
+def register_dataset(dataset_class):
+    dataset_hub[dataset_class.__name__.lower()] = dataset_class
+    return dataset_class

models/__init__.py

@@ -1,5 +1,4 @@
 import os, torch
-import segmentation_models_pytorch as smp
 
 from .adscnet import ADSCNet
 from .aglnet import AGLNet
@@ -38,51 +37,25 @@
 from .sqnet import SQNet
 from .stdc import STDC, LaplacianConv
 from .swiftnet import SwiftNet
-
-
-decoder_hub = {'deeplabv3':smp.DeepLabV3, 'deeplabv3p':smp.DeepLabV3Plus, 'fpn':smp.FPN,
-               'linknet':smp.Linknet, 'manet':smp.MAnet, 'pan':smp.PAN, 'pspnet':smp.PSPNet,
-               'unet':smp.Unet, 'unetpp':smp.UnetPlusPlus}
+from .model_registry import model_hub, aux_models, detail_head_models
 
 
 def get_model(config):
-    model_hub = {'adscnet':ADSCNet, 'aglnet':AGLNet, 'bisenetv1':BiSeNetv1, 
-                'bisenetv2':BiSeNetv2, 'canet':CANet, 'cfpnet':CFPNet, 
-                'cgnet':CGNet, 'contextnet':ContextNet, 'dabnet':DABNet, 
-                'ddrnet':DDRNet, 'dfanet':DFANet, 'edanet':EDANet, 
-                'enet':ENet, 'erfnet':ERFNet, 'esnet':ESNet, 
-                'espnet':ESPNet, 'espnetv2':ESPNetv2, 'fanet':FANet, 'farseenet':FarSeeNet, 
-                'fastscnn':FastSCNN, 'fddwnet':FDDWNet, 'fpenet':FPENet, 
-                'fssnet':FSSNet, 'icnet':ICNet, 'lednet':LEDNet,
-                'linknet':LinkNet, 'lite_hrnet':LiteHRNet, 'liteseg':LiteSeg, 'mininet':MiniNet, 
-                'mininetv2':MiniNetv2, 'ppliteseg':PPLiteSeg, 'regseg':RegSeg,
-                'segnet':SegNet, 'shelfnet':ShelfNet, 'sqnet':SQNet, 
-                'stdc':STDC, 'swiftnet':SwiftNet,}
-
-    # The following models currently support auxiliary heads
-    aux_models = ['bisenetv2', 'ddrnet', 'icnet']
-
-    # The following models currently support detail heads
-    detail_head_models = ['stdc']
-    
     if config.model == 'smp':   # Use segmentation models pytorch
-        if config.decoder not in decoder_hub:
-            raise ValueError(f"Unsupported decoder type: {config.decoder}")
+        from .smp_wrapper import get_smp_model
 
-        model = decoder_hub[config.decoder](encoder_name=config.encoder, 
-                                            encoder_weights=config.encoder_weights, 
-                                            in_channels=3, classes=config.num_class)
+        model = get_smp_model(config.encoder, config.decoder, config.encoder_weights, config.num_class)
 
     elif config.model in model_hub.keys():
-        if config.model in aux_models:
-            model = model_hub[config.model](num_class=config.num_class, use_aux=config.use_aux)
-        elif config.model in detail_head_models:
-            model = model_hub[config.model](num_class=config.num_class, use_detail_head=config.use_detail_head, use_aux=config.use_aux)
+        if config.model in aux_models:  # models support auxiliary heads
+            if config.model in detail_head_models:  # models support detail heads
+                model = model_hub[config.model](num_class=config.num_class, use_detail_head=config.use_detail_head, use_aux=config.use_aux)
+            else:
+                model = model_hub[config.model](num_class=config.num_class, use_aux=config.use_aux)
+
         else:
             if config.use_aux:
                 raise ValueError(f'Model {config.model} does not support auxiliary heads.\n')
-            if config.use_detail_head:
-                raise ValueError(f'Model {config.model} does not support detail heads.\n')
 
             model = model_hub[config.model](num_class=config.num_class)
 
@@ -92,16 +65,30 @@ def get_model(config):
     return model
 
 
+def list_available_models():
+    model_list = list(model_hub.keys())
+
+    try:
+        import segmentation_models_pytorch as smp
+        model_list.append('smp')
+    except:
+        pass
+
+    return model_list
+
+
 def get_teacher_model(config, device):
     if config.kd_training:
         if not os.path.isfile(config.teacher_ckpt):
-            raise ValueError(f'Could not find teacher checkpoint at path {config.teacher_ckpt}.')
+            raise ValueError(f'Could not find teacher checkpoint at path {config.teacher_ckpt}.')   
 
-        if config.teacher_decoder not in decoder_hub.keys():
-            raise ValueError(f"Unsupported teacher decoder type: {config.teacher_decoder}")      
+        if config.teacher_model == 'smp':
+            from .smp_wrapper import get_smp_model
 
-        model = decoder_hub[config.teacher_decoder](encoder_name=config.teacher_encoder, 
-                            encoder_weights=None, in_channels=3, classes=config.num_class)        
+            model = get_smp_model(config.teacher_encoder, config.teacher_decoder, None, config.num_class)
+
+        else:
+            raise NotImplementedError()
 
         teacher_ckpt = torch.load(config.teacher_ckpt, map_location=torch.device('cpu'))
         model.load_state_dict(teacher_ckpt['state_dict'])

models/adscnet.py

@@ -10,11 +10,13 @@
 import torch.nn as nn
 
 from .modules import conv1x1, ConvBNAct, DWConvBNAct, DeConvBNAct, Activation
+from .model_registry import register_model
 
 
+@register_model()
 class ADSCNet(nn.Module):
     def __init__(self, num_class=1, n_channel=3, act_type='relu6'):
-        super(ADSCNet, self).__init__()
+        super().__init__()
         self.conv0 = ConvBNAct(n_channel, 32, 3, 2, act_type=act_type, inplace=True)
         self.conv1 = ADSCModule(32, 1, act_type=act_type)
         self.conv2_4 = nn.Sequential(
@@ -54,7 +56,7 @@ def forward(self, x):
 
 class ADSCModule(nn.Module):
     def __init__(self, channels, stride, dilation=1, act_type='relu'):
-        super(ADSCModule, self).__init__()
+        super().__init__()
         assert stride in [1, 2], 'Unsupported stride type.\n'
         self.use_skip = stride == 1
         self.conv = nn.Sequential(
@@ -80,7 +82,7 @@ def forward(self, x):
 
 class DDCC(nn.Module):
     def __init__(self, channels, dilations, act_type):
-        super(DDCC, self).__init__()
+        super().__init__()
         assert len(dilations)==4, 'Length of dilations should be 4.\n'
         self.block1 = nn.Sequential(
                             nn.AvgPool2d(dilations[0], 1, dilations[0]//2),
@@ -109,17 +111,17 @@ def __init__(self, channels, dilations, act_type):
 
     def forward(self, x):
         x1 = self.block1(x)
-        
+
         x2 = torch.cat([x, x1], dim=1)
         x2 = self.block2(x2)
-        
+
         x3 = torch.cat([x, x1, x2], dim=1)
         x3 = self.block3(x3)
-        
+
         x4 = torch.cat([x, x1, x2, x3], dim=1)
         x4 = self.block4(x4)
-        
+
         x = torch.cat([x, x1, x2, x3, x4], dim=1)
         x = self.conv_last(x)
 
-        return x
+        return x