DataPipeline PoC

.gitignore

@@ -139,3 +139,4 @@ titanic.csv
 data_folder
 *.pt
 *.zip
+/data

flash/__init__.py

@@ -47,13 +47,13 @@
     # We are not importing the rest of the lightning during the build process, as it may not be compiled yet
 else:
 
-    from flash import tabular, text, vision
-    from flash.core import data, utils
+    from flash import data, tabular, text, vision
+    from flash.core import utils
     from flash.core.classification import ClassificationTask
-    from flash.core.data import DataModule
-    from flash.core.data.utils import download_data
     from flash.core.model import Task
     from flash.core.trainer import Trainer
+    from flash.data.data_module import DataModule
+    from flash.data.utils import download_data
 
     __all__ = [
         "Task",

flash/core/classification.py

@@ -15,23 +15,27 @@
 
 import torch
 
-from flash.core.data import TaskDataPipeline
 from flash.core.model import Task
+from flash.data.data_pipeline import Postprocess
 
 
-class ClassificationDataPipeline(TaskDataPipeline):
+class ClassificationDataPipeline:
+    pass
 
-    def before_uncollate(self, batch: Union[torch.Tensor, tuple]) -> torch.Tensor:
+
+class ClassificationPostprocess(Postprocess):
+
+    def per_batch_transform(self, batch: Union[torch.Tensor, tuple]) -> torch.Tensor:
         if isinstance(batch, tuple):
             batch = batch[0]
         return torch.softmax(batch, -1)
 
-    def after_uncollate(self, samples: Any) -> Any:
+    def per_sample_transform(self, samples: Any) -> Any:
         return torch.argmax(samples, -1).tolist()
 
 
 class ClassificationTask(Task):
 
-    @staticmethod
-    def default_pipeline() -> ClassificationDataPipeline:
-        return ClassificationDataPipeline()
+    @property
+    def postprocess(self):
+        return ClassificationPostprocess()

flash/core/finetuning.py

@@ -25,7 +25,7 @@ class NoFreeze(BaseFinetuning):
     def freeze_before_training(self, pl_module: pl.LightningModule) -> None:
         pass
 
-    def finetunning_function(
+    def finetune_function(
         self,
         pl_module: pl.LightningModule,
         epoch: int,
@@ -42,7 +42,7 @@ def __init__(self, attr_names: Union[str, List[str]] = "backbone", train_bn: boo
 
         FlashBaseFinetuning can be used to create a custom Flash Finetuning Callback.
 
-        Override ``finetunning_function`` to put your unfreeze logic.
+        Override ``finetune_function`` to put your unfreeze logic.
 
         Args:
             attr_names: Name(s) of the module attributes of the model to be frozen.
@@ -62,15 +62,15 @@ def freeze_using_attr_names(self, pl_module, attr_names: List[str], train_bn: bo
             attr = getattr(pl_module, attr_name, None)
             if attr is None or not isinstance(attr, nn.Module):
                 MisconfigurationException(f"Your model must have a {attr} attribute")
-            self.freeze(module=attr, train_bn=train_bn)
+            self.freeze(modules=attr, train_bn=train_bn)
 
-    def finetunning_function(self, pl_module: pl.LightningModule, epoch: int, optimizer: Optimizer, opt_idx: int):
+    def finetune_function(self, pl_module: pl.LightningModule, epoch: int, optimizer: Optimizer, opt_idx: int):
         pass
 
 
 class Freeze(FlashBaseFinetuning):
 
-    def finetunning_function(
+    def finetune_function(
         self,
         pl_module: pl.LightningModule,
         epoch: int,
@@ -86,7 +86,7 @@ def __init__(self, attr_names: Union[str, List[str]] = "backbone", train_bn: boo
         super().__init__(attr_names, train_bn)
         self.unfreeze_epoch = unfreeze_epoch
 
-    def finetunning_function(
+    def finetune_function(
         self,
         pl_module: pl.LightningModule,
         epoch: int,
@@ -116,7 +116,7 @@ def __init__(
 
         super().__init__(attr_names, train_bn)
 
-    def finetunning_function(
+    def finetune_function(
         self,
         pl_module: pl.LightningModule,
         epoch: int,