Renaming all instances of train_step to train_batch. Introducing validate|test|infer_batch to all existing models. Updated engine creation code to look for appropriate model methods. (#664)

Author: drewoldag

docs/data_flow.rst

@@ -24,7 +24,7 @@ The ``prepare_inputs`` function is responsible for taking in the output of the `
 Model Input and Output Pipeline
 -------------------------------
 
-The data is then either sent to ``pytorch`` ignite for training or to ``onnx`` for inference. In the case of training, the data is converted into a tensor and passed into the model's ``train_step`` function. The model processes the data and returns the output predictions. If the user wishes to perform data augmentations, these can be set up in the model's ``train_step`` function as well.
+The data is then either sent to ``pytorch`` ignite for training or to ``onnx`` for inference. In the case of training, the data is converted into a tensor and passed into the model's ``train_batch`` function. The model processes the data and returns the output predictions. If the user wishes to perform data augmentations, these can be set up in the model's ``train_batch`` function as well.
 
 In the ``onnx`` case, the data remains a numpy array throughout the model evaluation and to the result. Both paths result in the output being a numpy array.
 

docs/external_libraries.rst

@@ -43,7 +43,7 @@ Defining a Model
 ----------------
 
 Models must be written as a subclasses of ``torch.nn.Module``, use pytorch for computation, and
-be decorated with ``@hyrax_model``. Models must minimally define ``__init__``, ``forward``, and ``train_step``
+be decorated with ``@hyrax_model``. Models must minimally define ``__init__``, ``prepare_inputs``, ``infer_batch``, and ``train_batch``
 methods.
 
 In order to get the ``@hyrax_model`` decorator you can import it with ``from hyrax.models import hyrax_model``.
@@ -59,21 +59,21 @@ to allow your model class to adjust architecture or check that the provided data
 the first iterable axis of the numpy array.
 
 
-``forward(self, x)``
+``infer_batch(self, x)``
 ....................
 Hyrax calls this function, which evaluates your model on a single input ``x``. ``x`` is guaranteed to be a numpy array with
 the shape passed to ``__init__``.
 
-``forward()`` should return a numpy array that is the output of your model.
+``infer_batch(self, x)`` should return a numpy array that is the output of your model.
 
 
-``train_step(self, batch)``
+``train_batch(self, batch)``
 ...........................
 This is called several times every training epoch with a batch of input numpy arrays for your model, and is the
 inner training loop for your model. This is where you compute loss, perform back propagation, etc depending on
 how your model is trained.
 
-``train_step`` returns a dictionary with a "loss" key who's value is a list of loss values for the individual
+``train_batch`` returns a dictionary with a "loss" key who's value is a list of loss values for the individual
 items in the batch. This loss is logged to MLflow and tensorboard.
 
 Optional Methods

src/hyrax/models/hsc_autoencoder.py

@@ -43,15 +43,16 @@ def forward(self, x):
         decoded = self.decoder(encoded)
         return decoded
 
-    def train_step(self, batch):
+    def train_batch(self, batch):
         """
         This function contains the logic for a single training step. i.e. the
         contents of the inner loop of a ML training process.
 
         Parameters
         ----------
         batch : tuple
-            A tuple containing the two values the loss function
+            A tuple containing the input data for the current batch, possibly
+            with labels that are ignored.
 
         Returns
         -------
@@ -68,3 +69,74 @@ def train_step(self, batch):
         self.optimizer.step()
 
         return {"loss": loss.item()}
+
+    def validate_batch(self, batch):
+        """
+        This function contains the logic for a single validation step that will
+        process a single batch of data. i.e. the contents of the inner loop of a
+        ML validation process.
+
+        Parameters
+        ----------
+        batch : tuple
+            A tuple containing the input data for the current batch, possibly
+            with labels that are ignored.
+
+        Returns
+        -------
+        Current loss value : dict
+            Dictionary containing the loss value for the current batch.
+        """
+
+        data = batch[0]
+
+        decoded = self.forward(data)
+        loss = self.criterion(decoded, data)
+
+        return {"loss": loss.item()}
+
+    def test_batch(self, batch):
+        """
+        This function contains the logic for a single testing step that will
+        process a single batch of data. i.e. the contents of the inner loop of a
+        ML testing process. In this case, it is identical to `validate_batch`.
+
+        Parameters
+        ----------
+        batch : tuple
+            A tuple containing the input data for the current batch, possibly
+            with labels that are ignored.
+
+        Returns
+        -------
+        Current loss value : dict
+            Dictionary containing the loss value for the current batch.
+        """
+
+        data = batch[0]
+
+        decoded = self.forward(data)
+        loss = self.criterion(decoded, data)
+
+        return {"loss": loss.item()}
+
+    def infer_batch(self, batch):
+        """
+        This function contains the logic for a single inference step that will
+        process a single batch of data. i.e. the contents of the inner loop of a
+        ML inference process.
+
+        Parameters
+        ----------
+        batch : tuple
+            A tuple containing the input data for the current batch, possibly
+            with labels that are ignored.
+
+        Returns
+        -------
+        Reconstructed outputs : torch.Tensor
+            The reconstructed outputs from the autoencoder.
+        """
+
+        data = batch[0]
+        return self.forward(data)

src/hyrax/models/hsc_dcae.py

@@ -71,14 +71,15 @@ def forward(self, x):
 
         return x4
 
-    def train_step(self, batch):
+    def train_batch(self, batch):
         """This function contains the logic for a single training step. i.e. the
         contents of the inner loop of a ML training process.
 
         Parameters
         ----------
         batch : tuple
-            A tuple containing the two values the loss function
+            A tuple containing the input data for the current batch, possibly
+            with labels that are ignored.
 
         Returns
         -------
@@ -107,3 +108,93 @@ def train_step(self, batch):
         self.optimizer.step()
 
         return {"loss": loss.item()}
+
+    def validate_batch(self, batch):
+        """This function contains the logic for a single validation step that will
+        process a single batch of data. i.e. the contents of the inner loop of a
+        ML validation process.
+
+        Parameters
+        ----------
+        batch : tuple
+            A tuple containing the input data for the current batch, possibly
+            with labels that are ignored.
+
+        Returns
+        -------
+        Current loss value : dict
+            Dictionary containing the loss value for the current batch.
+        """
+
+        # Dropping labels if present
+        data = batch[0] if isinstance(batch, tuple) else batch
+
+        # Encoder with skip connections
+        x1 = self.activation(self.encoder1(data))
+        x2 = self.activation(self.encoder2(self.pool(x1)))
+        x3 = self.activation(self.encoder3(self.pool(x2)))
+        x4 = self.activation(self.encoder4(self.pool(x3)))
+
+        # Decoder with skip connections
+        x = self.activation(self.decoder4(x4) + x3)
+        x = self.activation(self.decoder3(x) + x2)
+        x = self.activation(self.decoder2(x) + x1)
+        decoded = self.final_activation(self.decoder1(x))
+
+        loss = self.criterion(decoded, data)
+
+        return {"loss": loss.item()}
+
+    def test_batch(self, batch):
+        """This function contains the logic for a single testing step that will
+        process a single batch of data. i.e. the contents of the inner loop of a
+        ML testing process. In this case, it is identical to `validate_batch`.
+
+        Parameters
+        ----------
+        batch : tuple
+            A tuple containing the input data for the current batch, possibly
+            with labels that are ignored.
+
+        Returns
+        -------
+        Current loss value : dict
+            Dictionary containing the loss value for the current batch.
+        """
+
+        # Dropping labels if present
+        data = batch[0] if isinstance(batch, tuple) else batch
+
+        # Encoder with skip connections
+        x1 = self.activation(self.encoder1(data))
+        x2 = self.activation(self.encoder2(self.pool(x1)))
+        x3 = self.activation(self.encoder3(self.pool(x2)))
+        x4 = self.activation(self.encoder4(self.pool(x3)))
+
+        # Decoder with skip connections
+        x = self.activation(self.decoder4(x4) + x3)
+        x = self.activation(self.decoder3(x) + x2)
+        x = self.activation(self.decoder2(x) + x1)
+        decoded = self.final_activation(self.decoder1(x))
+
+        loss = self.criterion(decoded, data)
+
+        return {"loss": loss.item()}
+
+    def infer_batch(self, batch):
+        """This function contains the logic for a single inference step that will
+        process a single batch of data. i.e. the contents of the inner loop of a
+        ML inference process.
+
+        Parameters
+        ----------
+        batch : tuple
+            A tuple containing the input data for the current batch, possibly
+            with labels that are ignored.
+
+        Returns
+        -------
+        Reconstructed outputs : torch.Tensor
+            The reconstructed outputs from the autoencoder.
+        """
+        return self.forward(batch)

src/hyrax/models/hyrax_autoencoder.py

@@ -20,7 +20,7 @@ class HyraxAutoencoder(nn.Module):
     This example model is taken from this
     `autoenocoder tutorial <https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/tutorial9/AE_CIFAR10.html>`_
 
-    The train function has been converted into train_step for use with pytorch-ignite.
+    The train function has been converted into train_batch for use with pytorch-ignite.
     """
 
     def __init__(self, config, data_sample=None):
@@ -109,14 +109,15 @@ def _eval_decoder(self, x):
     def forward(self, batch):
         return self._eval_encoder(batch)
 
-    def train_step(self, batch):
+    def train_batch(self, batch):
         """This function contains the logic for a single training step. i.e. the
         contents of the inner loop of a ML training process.
 
         Parameters
         ----------
         batch : tuple
-            A tuple containing the inputs and labels for the current batch.
+            A tuple containing the input data for the current batch, possibly
+            with labels that are ignored.
 
         Returns
         -------
@@ -136,6 +137,69 @@ def train_step(self, batch):
 
         return {"loss": loss.item()}
 
+    def validate_batch(self, batch):
+        """This function contains the logic for a single validation step that will
+        process a single batch of data. i.e. the contents of the inner loop of a
+        ML validation process.
+
+        Parameters
+        ----------
+        batch : tuple
+            A tuple containing the input data for the current batch, possibly
+            with labels that are ignored.
+
+        Returns
+        -------
+        Current loss value : dict
+            Dictionary containing the loss value for the current batch.
+        """
+        z = self._eval_encoder(batch)
+        x_hat = self._eval_decoder(z)
+        loss = F.mse_loss(batch, x_hat, reduction="none")
+        loss = loss.sum(dim=[1, 2, 3]).mean(dim=[0])
+
+        return {"loss": loss.item()}
+
+    def test_batch(self, batch):
+        """This function contains the logic for a single testing step that will
+        process a single batch of data. i.e. the contents of the inner loop of a
+        ML testing process. In this case, it is identical to `validate_batch`.
+
+        Parameters
+        ----------
+        batch : tuple
+            A tuple containing the input data for the current batch, possibly
+            with labels that are ignored.
+
+        Returns
+        -------
+        Current loss value : dict
+            Dictionary containing the loss value for the current batch.
+        """
+        z = self._eval_encoder(batch)
+        x_hat = self._eval_decoder(z)
+        loss = F.mse_loss(batch, x_hat, reduction="none")
+        loss = loss.sum(dim=[1, 2, 3]).mean(dim=[0])
+
+        return {"loss": loss.item()}
+
+    def infer_batch(self, batch):
+        """This function contains the logic for a single inference step. i.e. the
+        contents of the inner loop of a ML inference process.
+
+        Parameters
+        ----------
+        batch : tuple
+            A tuple containing the input data for the current batch, possibly
+            with labels that are ignored.
+
+        Returns
+        -------
+        Reconstructed inputs : torch.Tensor
+            The reconstructed inputs from the autoencoder.
+        """
+        return self.forward(batch)
+
     @staticmethod
     def prepare_inputs(data_dict) -> tuple:
         """This function converts structured data to the input tensor we need to run