diff --git a/kermt/util/parsing.py b/kermt/util/parsing.py index 71e6886..3e5a67a 100644 --- a/kermt/util/parsing.py +++ b/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, diff --git a/task/fingerprint.py b/task/fingerprint.py index bfbc8cf..d66ed29 100644 --- a/task/fingerprint.py +++ b/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 = [] diff --git a/task/kermttrainer.py b/task/kermttrainer.py index 5fae8c2..01b19ba 100644 --- a/task/kermttrainer.py +++ b/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) diff --git a/task/predict.py b/task/predict.py index caf71f4..1775743 100644 --- a/task/predict.py +++ b/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 = []