DataPipeline PoC

.gitignore

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

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 pre_uncollate(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 post_uncollate(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/data/datamodule.py

@@ -18,7 +18,7 @@
 import pytorch_lightning as pl
 from torch.utils.data import DataLoader, Dataset
 
-from flash.core.data.datapipeline import DataPipeline
+from flash.data.data_pipeline import DataPipeline, Postprocess, Preprocess
 
 
 class TaskDataPipeline(DataPipeline):
@@ -44,13 +44,15 @@ def __init__(
         train_ds: Optional[Dataset] = None,
         valid_ds: Optional[Dataset] = None,
         test_ds: Optional[Dataset] = None,
+        predict_ds: Optional[Dataset] = None,
         batch_size: int = 1,
         num_workers: Optional[int] = None,
     ):
         super().__init__()
         self._train_ds = train_ds
         self._valid_ds = valid_ds
         self._test_ds = test_ds
+        self._predict_ds = predict_ds
 
         if self._train_ds is not None:
             self.train_dataloader = self._train_dataloader
@@ -61,6 +63,9 @@ def __init__(
         if self._test_ds is not None:
             self.test_dataloader = self._test_dataloader
 
+        if self._predict_ds is not None:
+            self.predict_dataloader = self._predict_dataloader
+
         self.batch_size = batch_size
 
         # TODO: figure out best solution for setting num_workers
@@ -72,6 +77,8 @@ def __init__(
         self.num_workers = num_workers
 
         self._data_pipeline = None
+        self._preprocess = None
+        self._postprocess = None
 
     def _train_dataloader(self) -> DataLoader:
         return DataLoader(
@@ -80,7 +87,7 @@ def _train_dataloader(self) -> DataLoader:
             shuffle=True,
             num_workers=self.num_workers,
             pin_memory=True,
-            collate_fn=self.data_pipeline.collate_fn,
+            collate_fn=self.data_pipeline.worker_collate_fn,
             drop_last=True,
         )
 
@@ -90,7 +97,7 @@ def _val_dataloader(self) -> DataLoader:
             batch_size=self.batch_size,
             num_workers=self.num_workers,
             pin_memory=True,
-            collate_fn=self.data_pipeline.collate_fn,
+            collate_fn=self.data_pipeline.worker_collate_fn,
         )
 
     def _test_dataloader(self) -> DataLoader:
@@ -99,19 +106,44 @@ def _test_dataloader(self) -> DataLoader:
             batch_size=self.batch_size,
             num_workers=self.num_workers,
             pin_memory=True,
-            collate_fn=self.data_pipeline.collate_fn,
+            collate_fn=self.data_pipeline.worker_collate_fn,
+        )
+
+    def _predict_dataloader(self) -> DataLoader:
+        return DataLoader(
+            self._predict_ds,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            collate_fn=self.data_pipeline.worker_collate_fn,
         )
 
+    @property
+    def preprocess(self):
+        return self._preprocess
+
+    @preprocess.setter
+    def preprocess(self, preprocess: Preprocess) -> None:
+        self._preprocess = preprocess
+
+    @property
+    def postprocess(self):
+        return self._postprocess
+
+    @postprocess.setter
+    def postprocess(self, postprocess: Postprocess) -> None:
+        self._postprocess = postprocess
+
     @property
     def data_pipeline(self) -> DataPipeline:
         if self._data_pipeline is None:
-            self._data_pipeline = self.default_pipeline()
+            preprocess = self._preprocess
+            postprocess = self._postprocess
+            if preprocess is None and postprocess is None:
+                self._data_pipeline = self.default_pipeline()
+            return DataPipeline(preprocess, postprocess)
         return self._data_pipeline
 
     @data_pipeline.setter
     def data_pipeline(self, data_pipeline) -> None:
         self._data_pipeline = data_pipeline
-
-    @staticmethod
-    def default_pipeline() -> DataPipeline:
-        return TaskDataPipeline()

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,

flash/core/model.py

@@ -17,10 +17,12 @@
 
 import pytorch_lightning as pl
 import torch
+from pytorch_lightning import Trainer
 from torch import nn
 
-from flash.core.data import DataModule, DataPipeline
+from flash.core.data import DataModule
 from flash.core.utils import get_callable_dict
+from flash.data.data_pipeline import DataPipeline, Postprocess, Preprocess
 
 
 def predict_context(func: Callable) -> Callable:
@@ -31,13 +33,16 @@ def predict_context(func: Callable) -> Callable:
 
     @functools.wraps(func)
     def wrapper(self, *args, **kwargs) -> Any:
+        grad_enabled = torch.is_grad_enabled()
+        is_training = self.training
         self.eval()
         torch.set_grad_enabled(False)
 
         result = func(self, *args, **kwargs)
 
-        self.train()
-        torch.set_grad_enabled(True)
+        if is_training:
+            self.train()
+        torch.set_grad_enabled(grad_enabled)
         return result
 
     return wrapper
@@ -63,6 +68,8 @@ def __init__(
         learning_rate: float = 5e-5,
     ):
         super().__init__()
+        self._last_trainer_kwargs = {}
+
         if model is not None:
             self.model = model
         self.loss_fn = {} if loss_fn is None else get_callable_dict(loss_fn)
@@ -71,15 +78,18 @@ def __init__(
         self.learning_rate = learning_rate
         # TODO: should we save more? Bug on some regarding yaml if we save metrics
         self.save_hyperparameters("learning_rate", "optimizer")
+
         self._data_pipeline = None
+        self._preprocess = None
+        self._postprocess = None
 
     def step(self, batch: Any, batch_idx: int) -> Any:
         """
         The training/validation/test step. Override for custom behavior.
         """
         x, y = batch
         y_hat = self.forward(x)
-        output = {"y_hat": self.data_pipeline.before_uncollate(y_hat)}
+        output = {"y_hat": self.data_pipeline.pre_uncollate(y_hat)}
         losses = {name: l_fn(y_hat, y) for name, l_fn in self.loss_fn.items()}
         logs = {}
         for name, metric in self.metrics.items():
@@ -143,48 +153,73 @@ def predict(
             The post-processed model predictions
 
         """
-        # enable x to be a path to a folder
-        if isinstance(x, str):
-            files = os.listdir(x)
-            files = [os.path.join(x, y) for y in files]
-            x = files
-
         data_pipeline = data_pipeline or self.data_pipeline
-        batch = x if skip_collate_fn else data_pipeline.collate_fn(x)
-        batch_x, batch_y = batch if len(batch) == 2 and isinstance(batch, (list, tuple)) else (batch, None)
-        predictions = self.forward(batch_x)
-        output = data_pipeline.uncollate_fn(predictions)  # TODO: pass batch and x
-        return output
+        x = [x for x in data_pipeline._generate_auto_dataset(x)]
+        x = self.data_pipeline.worker_collate_fn(x)
+        #x = self.data_pipeline.device_collate_fn(x)
+        predictions = self.predict_step(x, batch_idx)
+        return data_pipeline.uncollate_fn(predictions)
+
+    def predict_step(self, batch, batch_idx):
+        return self(batch)
 
     def configure_optimizers(self) -> torch.optim.Optimizer:
         return self.optimizer_cls(filter(lambda p: p.requires_grad, self.parameters()), lr=self.learning_rate)
 
+    def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
+        self.data_pipeline = checkpoint["pipeline"]
+
+    def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
+        checkpoint["pipeline"] = self.data_pipeline
+
+    def configure_finetune_callback(self):
+        return []
+
+    def predict_step(self, batch, batch_idx):
+        return self(batch)
+
+    @property
+    def preprocess(self):
+        return self._preprocess
+
+    @preprocess.setter
+    def preprocess(self, preprocess: Preprocess) -> None:
+        data_pipeline = self.data_pipeline
+        self.data_pipeline = DataPipeline(preprocess, data_pipeline.postprocess)
+
+    @property
+    def postprocess(self):
+        return self._postprocess
+
+    @postprocess.setter
+    def postprocess(self, postprocess: Postprocess) -> None:
+        data_pipeline = self.data_pipeline
+        self.data_pipeline = DataPipeline(data_pipeline.preprocess, postprocess)
+
     @property
-    def data_pipeline(self) -> DataPipeline:
+    def data_pipeline(self) -> Optional[DataPipeline]:
         # we need to save the pipeline in case this class
         # is loaded from checkpoint and used to predict
-        if not self._data_pipeline:
-            try:
-                # datamodule pipeline takes priority
-                self._data_pipeline = self.trainer.datamodule.data_pipeline
-            except AttributeError:
-                self._data_pipeline = self.default_pipeline()
-        return self._data_pipeline
+        return self._get_pipeline("data_pipeline")
 
     @data_pipeline.setter
     def data_pipeline(self, data_pipeline: DataPipeline) -> None:
         self._data_pipeline = data_pipeline
+        if isinstance(data_pipeline, DataPipeline):
+            self._data_pipeline._attach_to_model(self)
 
-    @staticmethod
-    def default_pipeline() -> DataPipeline:
-        """Pipeline to use when there is no datamodule or it has not defined its pipeline"""
-        return DataModule.default_pipeline()
+    def _get_pipeline(self, pipeline_attr_name: str):
 
-    def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
-        self.data_pipeline = checkpoint["pipeline"]
+        if getattr(self, '_' + pipeline_attr_name) is not None:
+            return getattr(self, '_' + pipeline_attr_name)
 
-    def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
-        checkpoint["pipeline"] = self.data_pipeline
+        if self.datamodule is not None and hasattr(self, pipeline_attr_name):
+            return getattr(self.datamodule, pipeline_attr_name)
 
-    def configure_finetune_callback(self):
-        return []
+        if self.trainer is not None and hasattr(self.trainer, 'datamodule') and self.trainer.datamodule is not None:
+            if hasattr(self.trainer.datamodule,
+                       pipeline_attr_name) and getattr(self.trainer.datamodule, pipeline_attr_name):
+                data_pipeline = getattr(self.trainer.datamodule, pipeline_attr_name)
+                return DataPipeline(data_pipeline.preprocess, self.postprocess)
+
+        return None