feat(mean): add downsampled loss computation to improve efficiency

Author: torms3

deepem/loss/mean.py

@@ -78,6 +78,7 @@ def __init__(
         delta_d: float = 1.5,
         recompute_ext: bool = False,
         mask_background: bool = True,
+        loss_scale_factor: tuple[float, float, float] | None = None,
         **kwargs,
     ):
         super().__init__()
@@ -88,6 +89,7 @@ def __init__(
         self.delta_d = delta_d  # Distance (inter-cluster push force) hinge
         self.recompute_ext = recompute_ext
         self.mask_background = mask_background
+        self.loss_scale_factor = loss_scale_factor
 
     def forward(
         self,
@@ -104,6 +106,14 @@ def forward(
         """
         device = embd.device
 
+        # Downsample if enabled
+        if self.loss_scale_factor is not None:
+            embd = F.interpolate(embd, scale_factor=self.loss_scale_factor, mode='trilinear', align_corners=False)
+            trgt = F.interpolate(trgt, scale_factor=self.loss_scale_factor, mode='nearest')
+            mask = F.interpolate(mask, scale_factor=self.loss_scale_factor, mode='nearest')
+            if splt is not None:
+                splt = F.interpolate(splt, scale_factor=self.loss_scale_factor, mode='nearest')
+
         groups = None
         if self.recompute_ext:
             assert splt is not None

deepem/train/option.py

@@ -83,6 +83,7 @@ def initialize(self):
         self.parser.add_argument('--delta_d', type=float, default=1.5)
         self.parser.add_argument('--recompute_ext', action='store_true')
         self.parser.add_argument('--no_mask_background', action='store_true')
+        self.parser.add_argument('--loss_scale_factor', type=vec3f, default=None)
 
         # Optimizer
         self.parser.add_argument('--optim', default='Adam')
@@ -227,6 +228,7 @@ def parse(self):
         opt.metric_params['delta_d'] = opt.delta_d
         opt.metric_params['recompute_ext'] = opt.recompute_ext
         opt.metric_params['mask_background'] = not opt.no_mask_background
+        opt.metric_params['loss_scale_factor'] = opt.loss_scale_factor
 
         # Optimizer
         if opt.optim == 'Adam':