Update deepscreen/models/dti.py

deepscreen/models/dti.py

@@ -17,6 +17,8 @@ class DTILightningModule(LightningModule):
         model: a fully initialized instance of class torch.nn.Module
         metrics: a list of fully initialized instances of class torchmetrics.Metric
     """
+    extra_return_keys = ['ID1', 'X1', 'ID2', 'X2', 'N']
+
     def __init__(
             self,
             optimizer: optim.Optimizer,
@@ -58,13 +60,10 @@ class DTILightningModule(LightningModule):
             case 'predict':
                 dataloader = self.trainer.datamodule.predict_dataloader()
 
-
         # for key, value in dummy_batch.items():
         #     if isinstance(value, Tensor):
         #         dummy_batch[key] = value.to(self.device)
 
-
-
     def forward(self, batch):
         output = self.predictor(batch['X1^'], batch['X2^'])
         target = batch.get('Y')
@@ -94,13 +93,18 @@ class DTILightningModule(LightningModule):
         self.train_metrics(preds=preds, target=target, indexes=indexes.long())
         self.log_dict(self.train_metrics, on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)
 
-        return {
-            'N': batch['N'],
-            'ID1': batch['ID1'], 'X1': batch['X1'],
-            'ID2': batch['ID2'], 'X2': batch['X2'],
-            'Y^': preds, 'Y': target, 'loss': loss
+        return_dict = {
+            'Y^': preds,
+            'Y': target,
+            'loss': loss
         }
 
+        for key in self.extra_return_keys:
+            if key in batch:
+                return_dict[key] = batch[key]
+
+        return return_dict
+
     def on_train_epoch_end(self):
         pass
 
@@ -111,13 +115,18 @@ class DTILightningModule(LightningModule):
         self.val_metrics(preds=preds, target=target, indexes=indexes.long())
         self.log_dict(self.val_metrics, on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)
 
-        return {
-            'N': batch['N'],
-            'ID1': batch['ID1'], 'X1': batch['X1'],
-            'ID2': batch['ID2'], 'X2': batch['X2'],
-            'Y^': preds, 'Y': target, 'loss': loss
+        return_dict = {
+            'Y^': preds,
+            'Y': target,
+            'loss': loss
         }
 
+        for key in self.extra_return_keys:
+            if key in batch:
+                return_dict[key] = batch[key]
+
+        return return_dict
+
     def on_validation_epoch_end(self):
         pass
 
@@ -128,27 +137,34 @@ class DTILightningModule(LightningModule):
         self.test_metrics(preds=preds, target=target, indexes=indexes.long())
         self.log_dict(self.test_metrics, on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)
 
-        # return a dictionary for callbacks like BasePredictionWriter
-        return {
-            'N': batch['N'],
-            'ID1': batch['ID1'], 'X1': batch['X1'],
-            'ID2': batch['ID2'], 'X2': batch['X2'],
-            'Y^': preds, 'Y': target, 'loss': loss
+        return_dict = {
+            'Y^': preds,
+            'Y': target,
+            'loss': loss
         }
 
+        for key in self.extra_return_keys:
+            if key in batch:
+                return_dict[key] = batch[key]
+
+        return return_dict
+
     def on_test_epoch_end(self):
         pass
 
     def predict_step(self, batch, batch_idx, dataloader_idx=0):
         preds, _, _, _ = self.forward(batch)
         # return a dictionary for callbacks like BasePredictionWriter
-        return {
-            'N': batch['N'],
-            'ID1': batch['ID1'], 'X1': batch['X1'],
-            'ID2': batch['ID2'], 'X2': batch['X2'],
-            'Y^': preds
+        return_dict = {
+            'Y^': preds,
         }
 
+        for key in self.extra_return_keys:
+            if key in batch:
+                return_dict[key] = batch[key]
+
+        return return_dict
+
     def configure_optimizers(self):
         optimizers_config = {'optimizer': self.hparams.optimizer(params=self.parameters())}
         if self.hparams.get('scheduler'):