add ddpallreducemean

Author: CharlesPikachu

ssseg/modules/__init__.py

@@ -11,7 +11,6 @@
     BuildScheduler, BuildParamsConstructor, BuildNormalization, BuildActivation, BuildDropout, EMASegmentor
 )
 from .utils import (
-    initslurm, setrandomseed, touchdir, loadckpts, saveckpts, loadpretrainedweights, symlink, judgefileexist, postprocesspredgtpairs, ismainprocess, touchdirs,
-    LoggerHandleBuilder, BuildLoggerHandle, TrainingLoggingManager, BaseModuleBuilder, ConfigParser, EnvironmentCollector, SSSegInputStructure,
-    SSSegOutputStructure,
+    initslurm, setrandomseed, touchdir, loadckpts, saveckpts, loadpretrainedweights, symlink, judgefileexist, postprocesspredgtpairs, ismainprocess, touchdirs, ddpallreducemean,
+    LoggerHandleBuilder, TrainingLoggingManager, BaseModuleBuilder, ConfigParser, EnvironmentCollector, SSSegInputStructure, SSSegOutputStructure, BuildLoggerHandle,
 )

ssseg/modules/models/segmentors/ocrnet/ocrnet.py

@@ -56,7 +56,7 @@ def forward(self, data_meta):
         # feed to bottleneck
         feats = self.bottleneck(backbone_outputs[-1])
         # feed to ocr module
-        context = self.spatial_gather_module(feats, seg_logits_aux)
+        context = self.spatial_gather_module(feats, F.interpolate(seg_logits_aux, size=feats.shape[-2:], mode='bilinear', align_corners=self.align_corners))
         feats = self.object_context_block(feats, context)
         # feed to decoder
         seg_logits = self.decoder(feats)

ssseg/modules/utils/__init__.py

@@ -4,6 +4,6 @@
 from .configparser import ConfigParser
 from .modulebuilder import BaseModuleBuilder
 from .datastructure import SSSegInputStructure, SSSegOutputStructure
-from .misc import setrandomseed, postprocesspredgtpairs, ismainprocess
 from .logger import LoggerHandleBuilder, BuildLoggerHandle, TrainingLoggingManager
+from .misc import setrandomseed, postprocesspredgtpairs, ismainprocess, ddpallreducemean
 from .io import touchdir, loadckpts, saveckpts, loadpretrainedweights, symlink, judgefileexist, touchdirs

ssseg/modules/utils/misc.py

@@ -13,6 +13,14 @@
 from .io import touchdirs
 
 
+'''ddpallreducemean'''
+def ddpallreducemean(tensor: torch.Tensor) -> torch.Tensor:
+    world_size = dist.get_world_size()
+    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+    tensor /= world_size
+    return tensor
+
+
 '''setrandomseed'''
 def setrandomseed(seed):
     random.seed(seed)