Update training, validation, and testing steps to match the shared_step() design pattern

Author: amorehead

project/utils/deepinteract_modules.py

@@ -1761,7 +1761,11 @@ def training_step(self, batch, batch_idx):
         """Lightning calls this inside the training loop."""
         # Separate training batch from validation batch (with the latter being used for visualizations)
         train_batch, val_batch = batch['train_batch'], batch['val_batch']
-        graph1, graph2, examples_list, filepaths = train_batch[0], train_batch[1], train_batch[2], train_batch[3]
+
+        # Unpack list of input training complexes - Assume batch_size=1
+        train_batch = train_batch[0]
+        graph1, graph2 = train_batch['graph1'], train_batch['graph2']
+        examples_list, filepaths = [train_batch['examples']], [train_batch['filepath']]
 
         # Forward propagate with network layers
         logits_list, _, _, _, _ = self.shared_step(graph1, graph2)  # The forward method must be named something new
@@ -1832,10 +1836,13 @@ def training_step(self, batch, batch_idx):
                 # Val Sample
                 # ------------
                 # Make a forward pass through the network for a held-out validation complex for visualization
-                val_graph1, val_graph2, val_examples_list = val_batch[0], val_batch[1], val_batch[2][0]
+                val_batch = val_batch[0]  # Unpack list of input validation complexes - Assume batch_size=1
+                val_graph1, val_graph2 = val_batch['graph1'], val_batch['graph2']
+                val_examples_list = val_batch['examples']
 
                 # Forward propagate with network layers without accumulating any gradients
-                val_logits, _, _, _, _ = self.shared_step(val_graph1, val_graph2)[0].squeeze()
+                val_logits_list, _, _, _, _ = self.shared_step(val_graph1, val_graph2)
+                val_logits = val_logits_list[0].squeeze()
                 val_len_1, val_len_2 = val_logits.shape[1:]
 
                 # Construct the predicted M x N interaction tensor and its corresponding labels
@@ -1916,7 +1923,9 @@ def training_epoch_end(self, outputs: pl.utilities.types.EPOCH_OUTPUT) -> None:
     def validation_step(self, batch, batch_idx):
         """Lightning calls this inside the validation loop."""
         # Make a forward pass through the network for a batch of protein complexes
-        graph1, graph2, examples_list, filepaths = batch[0], batch[1], batch[2], batch[3]
+        batch = batch[0]  # Unpack list of input complexes - Assume batch_size=1
+        graph1, graph2 = batch['graph1'], batch['graph2']
+        examples_list, filepaths = [batch['examples']], [batch['filepath']]
 
         # Forward propagate with network layers
         logits_list, _, _, _, _ = self.shared_step(graph1, graph2)
@@ -2018,8 +2027,10 @@ def validation_epoch_end(self, outputs: pl.utilities.types.EPOCH_OUTPUT) -> None
 
     def test_step(self, batch, batch_idx):
         """Lightning calls this inside the testing loop."""
-        # Make a forward pass through the network for a batch of protein complexes (batch_size=1)
-        graph1, graph2, examples_list, filepaths = batch[0], batch[1], batch[2], batch[3]
+        # Make a forward pass through the network for a batch of protein complexes
+        batch = batch[0]  # Unpack list of input complexes - Assume batch_size=1
+        graph1, graph2 = batch['graph1'], batch['graph2']
+        examples_list, filepaths = [batch['examples']], [batch['filepath']]
 
         # Forward propagate with network layers
         logits_list, _, _, _, _ = self.shared_step(graph1, graph2)