make some update

Author: CharlesPikachu

csseg/configs/rcil/base_cfg.py

@@ -32,6 +32,7 @@
             'loss_seg': {'MIBUnbiasedCrossEntropyLoss': {'scale_factor': 1.0, 'reduction': 'mean', 'ignore_index': 255}}
         },
         'distillation': {'scale_factor': 1.0, 'spp_scales': [4, 8, 12, 16, 20, 24]},
+        'distillation_mib': {'scale_factor': 100, 'alpha': 1.0},
     }
 }
 # RUNNER_CFG

csseg/modules/models/decoders/rcilaspphead.py

@@ -57,8 +57,10 @@ def __init__(self, in_channels, feats_channels, out_channels, dilations, pooling
             nn.LeakyReLU(0.01),
         )
         # initialize parameters
-        assert norm_cfg['activation'] == 'identity'
-        self.initparams(actname2torchactname(act_cfg['type']), act_cfg.get('negative_slope'))
+        if hasattr(self.bottleneck_bn[0], 'activation'):
+            self.initparams(self.bottleneck_bn[0].activation, self.bottleneck_bn[0].activation_param)
+        else:
+            self.initparams(actname2torchactname(act_cfg['type']), act_cfg.get('negative_slope'))
     '''initparams'''
     def initparams(self, nonlinearity, param=None):
         gain = nn.init.calculate_gain(nonlinearity, param)

csseg/modules/models/encoders/resnetcaf.py

@@ -0,0 +1,7 @@
+'''
+Function:
+    Implementation of ResNetUCD
+Author:
+    Zhenchao Jin
+'''
+from .resnetmib import ResNetMIB as ResNetUCD

csseg/modules/models/segmentors/caf.py

@@ -1,6 +1,6 @@
 '''
 Function:
-    Implementation of "Class Similarity Weighted Knowledge Distillation for Continual Semantic Segmentation"
+    Implementation of "Continual attentive fusion for incremental learning in semantic segmentation"
 Author:
     Zhenchao Jin
 '''

csseg/modules/models/segmentors/ewf.py

@@ -12,6 +12,7 @@
 import torch.nn.functional as F
 import torch.distributed as dist
 from apex import amp
+from .mib import MIBRunner
 from .base import BaseRunner
 
 
@@ -153,7 +154,7 @@ def train(self, cur_epoch):
                 seg_targets=seg_targets, 
                 losses_cfgs=seg_losses_cfgs,
             )
-            # --calculate distillation losses
+            # --calculate pod distillation losses
             pod_total_loss, pod_losses_log_dict = 0, {}
             if self.history_segmentor is not None:
                 distillation_feats = outputs['distillation_feats']
@@ -165,8 +166,16 @@ def train(self, cur_epoch):
                     dataset_type=self.runner_cfg['dataset_cfg']['type'],
                     **losses_cfgs['distillation']
                 )
-            # --merge two losses
-            loss_total = pod_total_loss + seg_total_loss
+            # --calculate mib distillation losses
+            kd_total_loss, kd_losses_log_dict = 0, {}
+            if self.history_segmentor is not None:
+                kd_total_loss, kd_losses_log_dict = MIBRunner.featuresdistillation(
+                    history_distillation_feats=F.interpolate(history_outputs['seg_logits'], size=images.shape[2:], mode="bilinear", align_corners=self.segmentor.module.align_corners), 
+                    distillation_feats=F.interpolate(outputs['seg_logits'], size=images.shape[2:], mode="bilinear", align_corners=self.segmentor.module.align_corners),
+                    **losses_cfgs['distillation_mib']
+                )
+            # --merge three losses
+            loss_total = pod_total_loss + kd_total_loss + seg_total_loss
             # --perform back propagation
             with amp.scale_loss(loss_total, self.optimizer) as scaled_loss_total:
                 scaled_loss_total.backward()
@@ -175,6 +184,7 @@ def train(self, cur_epoch):
             self.scheduler.zerograd()
             # --logging training loss info
             seg_losses_log_dict.update(pod_losses_log_dict)
+            seg_losses_log_dict.update(kd_losses_log_dict)
             seg_losses_log_dict.pop('loss_total')
             seg_losses_log_dict['loss_total'] = loss_total.item()
             losses_log_dict = self.loggingtraininginfo(seg_losses_log_dict, losses_log_dict, init_losses_log_dict)

csseg/modules/runners/caf.py

@@ -1,6 +1,6 @@
 '''
 Function:
-    Implementation of "Continual attentive fusion for incremental learning in semantic segmentation"
+    Implementation of "Class Similarity Weighted Knowledge Distillation for Continual Semantic Segmentation"
 Author:
     Zhenchao Jin
 '''