Repository: Lightning-Universe/lightning-flash

Issue:
Output dictionary keys hardcoded in model.py should be in DefaultDataKeys
https://github.com/PyTorchLightning/lightning-flash/blob/143112d01542083214d8d6b1bfba54688ac9cf3e/flash/core/model.py#L353-L355

It seems to me that hardcoding "loss" "y_hat" and "y" in here isn't the best practice and encourages everyone who needs those values in callbacks to either hardcode or store the key values in their own variables. For whatever reason, if the keys were to change, custom callbacks would break easily. 

I propose moving the output keys into the `DefaultDataKeys` for this use case. Naming them something with a prefix so as not to be confused with other input data keys. Alternatively, make a different enum class for these to go in and be used.



Discussion:
@dlangerm Thanks for the suggestion! Definitely good to get rid of those magic strings. I think either of your proposed solutions would be fine. Perhaps a new enum is best as it could reside in the same file as the strings for readability. Would you be interested in making a PR to do that?

Thanks :smiley:

Resolving patch:
diff --git a/flash/core/model.py b/flash/core/model.py
--- a/flash/core/model.py
+++ b/flash/core/model.py
@@ -26,6 +26,7 @@
 from pytorch_lightning.callbacks import Callback
 from pytorch_lightning.trainer.states import RunningStage
 from pytorch_lightning.utilities import rank_zero_warn
+from pytorch_lightning.utilities.enums import LightningEnum
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from torch import nn
 from torch.optim.lr_scheduler import _LRScheduler
@@ -267,6 +268,19 @@ def __new__(mcs, *args, **kwargs):
         return result
 
 
+class OutputKeys(LightningEnum):
+    """The ``OutputKeys`` enum contains the keys that are used internally by the ``Task`` when handling outputs."""
+
+    OUTPUT = "y_hat"
+    TARGET = "y"
+    LOGS = "logs"
+    LOSS = "loss"
+
+    # TODO: Create a FlashEnum class???
+    def __hash__(self) -> int:
+        return hash(self.value)
+
+
 class Task(DatasetProcessor, ModuleWrapperBase, LightningModule, metaclass=CheckDependenciesMeta):
     """A general Task.
 
@@ -331,11 +345,11 @@ def step(self, batch: Any, batch_idx: int, metrics: nn.ModuleDict) -> Any:
         x, y = batch
         y_hat = self(x)
         y, y_hat = self.apply_filtering(y, y_hat)
-        output = {"y_hat": y_hat}
-        y_hat = self.to_loss_format(output["y_hat"])
+        output = {OutputKeys.OUTPUT: y_hat}
+        y_hat = self.to_loss_format(output[OutputKeys.OUTPUT])
         losses = {name: l_fn(y_hat, y) for name, l_fn in self.loss_fn.items()}
 
-        y_hat = self.to_metrics_format(output["y_hat"])
+        y_hat = self.to_metrics_format(output[OutputKeys.OUTPUT])
 
         logs = {}
 
@@ -350,9 +364,9 @@ def step(self, batch: Any, batch_idx: int, metrics: nn.ModuleDict) -> Any:
             logs["total_loss"] = sum(losses.values())
             return logs["total_loss"], logs
 
-        output["loss"] = self.compute_loss(losses)
-        output["logs"] = self.compute_logs(logs, losses)
-        output["y"] = y
+        output[OutputKeys.LOSS] = self.compute_loss(losses)
+        output[OutputKeys.LOGS] = self.compute_logs(logs, losses)
+        output[OutputKeys.TARGET] = y
         return output
 
     def compute_loss(self, losses: Dict[str, torch.Tensor]) -> torch.Tensor:
@@ -380,16 +394,31 @@ def forward(self, x: Any) -> Any:
 
     def training_step(self, batch: Any, batch_idx: int) -> Any:
         output = self.step(batch, batch_idx, self.train_metrics)
-        self.log_dict({f"train_{k}": v for k, v in output["logs"].items()}, on_step=True, on_epoch=True, prog_bar=True)
-        return output["loss"]
+        self.log_dict(
+            {f"train_{k}": v for k, v in output[OutputKeys.LOGS].items()},
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+        )
+        return output[OutputKeys.LOSS]
 
     def validation_step(self, batch: Any, batch_idx: int) -> None:
         output = self.step(batch, batch_idx, self.val_metrics)
-        self.log_dict({f"val_{k}": v for k, v in output["logs"].items()}, on_step=False, on_epoch=True, prog_bar=True)
+        self.log_dict(
+            {f"val_{k}": v for k, v in output[OutputKeys.LOGS].items()},
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+        )
 
     def test_step(self, batch: Any, batch_idx: int) -> None:
         output = self.step(batch, batch_idx, self.val_metrics)
-        self.log_dict({f"test_{k}": v for k, v in output["logs"].items()}, on_step=False, on_epoch=True, prog_bar=True)
+        self.log_dict(
+            {f"test_{k}": v for k, v in output[OutputKeys.LOGS].items()},
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+        )
 
     @predict_context
     def predict(


Test patch: