fix(predict+pretrain): rdkit2D predict arg, issue #22, and atom/bond-mode DDP + float-loss fixes

kermt/util/parsing.py

@@ -86,6 +86,9 @@ def add_predict_args(parser: ArgumentParser):
     parser.add_argument('--features_generator', type=str, nargs='*',
                         choices=get_available_features_generators(),
                         help='Method of generating additional features')
+    parser.add_argument('--rdkit2D_normalization_type', type=str, choices=("fast", "best", "descriptastorus"), default='fast',
+                        help='Type of normalization for rdkit2D features. Choices: fast, best, descriptastorus. '
+                             'Must match the value used at finetune time (it is baked into the features).')
     parser.add_argument('--features_path', type=str, nargs='*',
                         help='Path to features to use in FNN (instead of features_generator)')
     parser.add_argument('--use_cuikmolmaker_featurization', action='store_true', default=False,

task/fingerprint.py

@@ -37,6 +37,7 @@
 """
 The fingerprint generation function.
 """
+import copy
 from argparse import Namespace
 from logging import Logger
 from typing import List
@@ -63,6 +64,9 @@ def do_generate(model: nn.Module,
     :return: A list of fingerprints.
     """
     model.eval()
+    # Disable bond dropout locally without mutating the shared args (same pattern
+    # as predict(); see issue #22).
+    args = copy.copy(args)
     args.bond_drop_rate = 0
     preds = []
 

task/kermttrainer.py

@@ -215,7 +215,13 @@ def __init__(self,
 
         self.model.to(self.gpu_id)
 
-        self.model = DDP(self.model, device_ids=[gpu_id])
+        # find_unused_parameters is required when embedding_output_type != 'both':
+        # in atom/bond mode the encoder produces only one aggregation level, so the
+        # opposite branch's FFNs and the bond/atom vocab + FG heads receive no gradient.
+        # 'both' exercises every branch, so keep the flag off there to avoid DDP overhead.
+        # (static_graph is not an option: dynamic-depth sampling varies the graph per step.)
+        self.model = DDP(self.model, device_ids=[gpu_id],
+                         find_unused_parameters=(self.args.embedding_output_type != 'both'))
 
         if self.args.tensorboard:
             self.writer = SummaryWriter(self.args.save_dir)
@@ -270,9 +276,9 @@ def validation(self, max_val_batches: int) -> List:
             loss_sum += loss.item()
             iter_count += self.args.batch_size
 
-            av_loss_sum += av_loss.item()
-            bv_loss_sum += bv_loss.item()
-            fg_loss_sum += fg_loss.item()
+            av_loss_sum += av_loss.item() if not isinstance(av_loss, float) else av_loss
+            bv_loss_sum += bv_loss.item() if not isinstance(bv_loss, float) else bv_loss
+            fg_loss_sum += fg_loss.item() if not isinstance(fg_loss, float) else fg_loss
             av_dist_loss_sum += av_dist_loss.item() if type(av_dist_loss) != float else av_dist_loss
             bv_dist_loss_sum += bv_dist_loss.item() if type(bv_dist_loss) != float else bv_dist_loss
             fg_dist_loss_sum += fg_dist_loss.item() if type(fg_dist_loss) != float else fg_dist_loss
@@ -395,13 +401,13 @@ def iter(self, epoch, train=True) -> List:
                 self.scheduler.step()
             else:
                 # For eval model, only consider the loss of three task.
-                cum_loss_sum += av_loss.item()
-                cum_loss_sum += bv_loss.item()
-                cum_loss_sum += fg_loss.item()
+                cum_loss_sum += av_loss.item() if not isinstance(av_loss, float) else av_loss
+                cum_loss_sum += bv_loss.item() if not isinstance(bv_loss, float) else bv_loss
+                cum_loss_sum += fg_loss.item() if not isinstance(fg_loss, float) else fg_loss
 
-            av_loss_sum += av_loss.item()
-            bv_loss_sum += bv_loss.item()
-            fg_loss_sum += fg_loss.item()
+            av_loss_sum += av_loss.item() if not isinstance(av_loss, float) else av_loss
+            bv_loss_sum += bv_loss.item() if not isinstance(bv_loss, float) else bv_loss
+            fg_loss_sum += fg_loss.item() if not isinstance(fg_loss, float) else fg_loss
             av_dist_loss_sum += av_dist_loss.item() if type(av_dist_loss) != float else av_dist_loss
             bv_dist_loss_sum += bv_dist_loss.item() if type(bv_dist_loss) != float else bv_dist_loss
             fg_dist_loss_sum += fg_dist_loss.item() if type(fg_dist_loss) != float else fg_dist_loss
@@ -432,12 +438,12 @@ def iter(self, epoch, train=True) -> List:
             if self.gpu_id == 0 and self.n_steps % train_log_interval == 0:
                 train_metrics = {
                     'train/loss': loss.item(),
-                    'train/av_loss': av_loss.item(),
-                    'train/bv_loss': bv_loss.item(),
-                    'train/fg_loss': fg_loss.item(),
-                    'train/av_dist_loss': av_dist_loss.item(),
-                    'train/bv_dist_loss': bv_dist_loss.item(),
-                    'train/fg_dist_loss': fg_dist_loss.item(),
+                    'train/av_loss': av_loss.item() if not isinstance(av_loss, float) else av_loss,
+                    'train/bv_loss': bv_loss.item() if not isinstance(bv_loss, float) else bv_loss,
+                    'train/fg_loss': fg_loss.item() if not isinstance(fg_loss, float) else fg_loss,
+                    'train/av_dist_loss': av_dist_loss.item() if not isinstance(av_dist_loss, float) else av_dist_loss,
+                    'train/bv_dist_loss': bv_dist_loss.item() if not isinstance(bv_dist_loss, float) else bv_dist_loss,
+                    'train/fg_dist_loss': fg_dist_loss.item() if not isinstance(fg_dist_loss, float) else fg_dist_loss,
                     'train/lr': self.scheduler.get_lr()[0],
                     'train/epoch': epoch,
                     'train/batch_idx': ibatch + self.batch_idx_offset,
@@ -564,7 +570,13 @@ def __init__(self,
         self.n_iter = 0
 
         self.model.to(self.gpu_id)
-        self.model = DDP(self.model, device_ids=[gpu_id])
+        # find_unused_parameters is required when embedding_output_type != 'both':
+        # in atom/bond mode the encoder produces only one aggregation level, so the
+        # opposite branch's FFNs and the bond/atom vocab + FG heads receive no gradient.
+        # 'both' exercises every branch, so keep the flag off there to avoid DDP overhead.
+        # (static_graph is not an option: dynamic-depth sampling varies the graph per step.)
+        self.model = DDP(self.model, device_ids=[gpu_id],
+                         find_unused_parameters=(self.args.embedding_output_type != 'both'))
 
         if self.args.tensorboard:
             self.writer = SummaryWriter(self.args.save_dir)
@@ -871,7 +883,13 @@ def __init__(self,
         self.n_iter = 0
 
         self.model.to(self.gpu_id)
-        self.model = DDP(self.model, device_ids=[gpu_id])
+        # find_unused_parameters is required when embedding_output_type != 'both':
+        # in atom/bond mode the encoder produces only one aggregation level, so the
+        # opposite branch's FFNs and the bond/atom vocab + FG heads receive no gradient.
+        # 'both' exercises every branch, so keep the flag off there to avoid DDP overhead.
+        # (static_graph is not an option: dynamic-depth sampling varies the graph per step.)
+        self.model = DDP(self.model, device_ids=[gpu_id],
+                         find_unused_parameters=(self.args.embedding_output_type != 'both'))
 
         if self.args.tensorboard:
             self.writer = SummaryWriter(self.args.save_dir)

task/predict.py

@@ -37,6 +37,7 @@
 """
 The predict function using the finetuned model to make the prediction. .
 """
+import copy
 from argparse import Namespace
 from typing import List
 
@@ -75,6 +76,13 @@ def predict(model: nn.Module,
     """
     # debug = logger.debug if logger is not None else print
     model.eval()
+    # Evaluation must not apply bond dropout, but do NOT mutate the shared args:
+    # training and eval reuse one args instance and graphs are rebuilt every epoch
+    # (caching is off by default), so setting args.bond_drop_rate = 0 here would
+    # permanently disable bond dropout for all later training epochs -- and under
+    # DDP, where only rank 0 evaluates, desync augmentation across ranks. Use a
+    # shallow copy so the override is local to this call. See issue #22.
+    args = copy.copy(args)
     args.bond_drop_rate = 0
     preds = []