Merge pull request IntelLabs#185 from melo-gonzo/sam-multi-pipeline-fix

SAM Multi Task/Data Pipeline Fix

matsciml/lightning/callbacks.py

@@ -8,6 +8,7 @@
 from logging import DEBUG, getLogger
 from pathlib import Path
 from time import time
+from copy import copy
 from typing import Any, Callable, Dict, Iterator, Optional
 
 import numpy as np
@@ -753,6 +754,21 @@ def on_train_batch_start(
         self.batch = batch
         self.batch_idx = batch_idx
 
+    def extract_optimizer_specific_loss(self, trainer, optimizer, loss):
+        optimizer_names = copy(trainer.model.optimizer_names)
+        opt_idx = [opt == optimizer for opt in trainer.optimizers].index(True)
+        loss_keys = optimizer_names[opt_idx]
+        if loss_keys == ("Global", "Encoder"):
+            optimizer_names.pop(opt_idx)
+            global_loss = 0
+            for dataset, task in optimizer_names:
+                global_loss += loss[dataset][task]["loss"]
+            return {"loss": global_loss}
+        else:
+            for key in loss_keys:
+                loss = loss[key]
+        return loss
+
     def on_before_optimizer_step(
         self,
         trainer: Trainer,
@@ -763,6 +779,8 @@ def on_before_optimizer_step(
             org_weights = self._first_step(optimizer)
         with torch.enable_grad():
             loss = task._compute_losses(self.batch)
+            if len(trainer.optimizers) > 1:
+                loss = self.extract_optimizer_specific_loss(trainer, optimizer, loss)
             loss = self._get_loss(loss)
             if torch.isfinite(loss):
                 trainer.strategy.backward(loss, optimizer=optimizer)

matsciml/lightning/tests/test_sam.py

@@ -365,5 +365,5 @@ def test_multitask_sam():
         ("IS2REDataset", is2re),
         ("S2EFDataset", s2ef),
     )
-    trainer = pl.Trainer(fast_dev_run=10)
+    trainer = pl.Trainer(fast_dev_run=10, callbacks=SAM())
     trainer.fit(task, datamodule=dm)