[PoC] Add MetaLearning support through learn2learn

.azure-pipelines/gpu-tests.yml

@@ -59,6 +59,11 @@ jobs:
         python -m coverage run --source flash -m pytest flash tests/examples/test_scripts.py -v --junitxml=$(Build.StagingDirectory)/test-results.xml --durations=30
       displayName: 'Testing'
 
+    - bash: |
+        pip install git+https://github.com/tchaton/learn2learn@flash
+        bash tests/special_tests.sh
+      displayName: 'Testing: special'
+
     - bash: |
         python -m coverage report
         python -m coverage xml

flash/core/adapter.py

@@ -14,6 +14,7 @@
 from abc import abstractmethod
 from typing import Any, Callable, Optional
 
+import torch.jit
 from torch import nn
 from torch.utils.data import DataLoader, Sampler
 
@@ -59,6 +60,10 @@ def test_epoch_end(self, outputs) -> None:
         pass
 
 
+def identity_collate_fn(x):
+    return x
+
+
 class AdapterTask(Task):
     """The ``AdapterTask`` is a :class:`~flash.core.model.Task` which wraps an :class:`~flash.core.adapter.Adapter`
     and forwards all of the hooks.
@@ -73,11 +78,12 @@ def __init__(self, adapter: Adapter, **kwargs):
 
         self.adapter = adapter
 
+    @torch.jit.unused
     @property
     def backbone(self) -> nn.Module:
         return self.adapter.backbone
 
-    def forward(self, x: Any) -> Any:
+    def forward(self, x: torch.Tensor) -> Any:
         return self.adapter.forward(x)
 
     def training_step(self, batch: Any, batch_idx: int) -> Any:
@@ -104,6 +110,7 @@ def test_epoch_end(self, outputs) -> None:
     def process_train_dataset(
         self,
         dataset: BaseAutoDataset,
+        trainer: flash.Trainer,
         batch_size: int,
         num_workers: int,
         pin_memory: bool,
@@ -113,12 +120,13 @@ def process_train_dataset(
         sampler: Optional[Sampler] = None,
     ) -> DataLoader:
         return self.adapter.process_train_dataset(
-            dataset, batch_size, num_workers, pin_memory, collate_fn, shuffle, drop_last, sampler
+            dataset, trainer, batch_size, num_workers, pin_memory, collate_fn, shuffle, drop_last, sampler
         )
 
     def process_val_dataset(
         self,
         dataset: BaseAutoDataset,
+        trainer: flash.Trainer,
         batch_size: int,
         num_workers: int,
         pin_memory: bool,
@@ -128,12 +136,13 @@ def process_val_dataset(
         sampler: Optional[Sampler] = None,
     ) -> DataLoader:
         return self.adapter.process_val_dataset(
-            dataset, batch_size, num_workers, pin_memory, collate_fn, shuffle, drop_last, sampler
+            dataset, trainer, batch_size, num_workers, pin_memory, collate_fn, shuffle, drop_last, sampler
         )
 
     def process_test_dataset(
         self,
         dataset: BaseAutoDataset,
+        trainer: flash.Trainer,
         batch_size: int,
         num_workers: int,
         pin_memory: bool,
@@ -143,7 +152,7 @@ def process_test_dataset(
         sampler: Optional[Sampler] = None,
     ) -> DataLoader:
         return self.adapter.process_test_dataset(
-            dataset, batch_size, num_workers, pin_memory, collate_fn, shuffle, drop_last, sampler
+            dataset, trainer, batch_size, num_workers, pin_memory, collate_fn, shuffle, drop_last, sampler
         )
 
     def process_predict_dataset(
@@ -152,11 +161,18 @@ def process_predict_dataset(
         batch_size: int = 1,
         num_workers: int = 0,
         pin_memory: bool = False,
-        collate_fn: Callable = lambda x: x,
+        collate_fn: Callable = identity_collate_fn,
         shuffle: bool = False,
         drop_last: bool = True,
         sampler: Optional[Sampler] = None,
     ) -> DataLoader:
         return self.adapter.process_predict_dataset(
-            dataset, batch_size, num_workers, pin_memory, collate_fn, shuffle, drop_last, sampler
+            dataset,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            pin_memory=pin_memory,
+            collate_fn=collate_fn,
+            shuffle=shuffle,
+            drop_last=drop_last,
+            sampler=sampler,
         )

flash/core/classification.py

@@ -18,6 +18,7 @@
 import torchmetrics
 from pytorch_lightning.utilities import rank_zero_warn
 
+from flash.core.adapter import AdapterTask
 from flash.core.data.data_source import DefaultDataKeys, LabelsState
 from flash.core.data.process import Serializer
 from flash.core.model import Task
@@ -37,7 +38,29 @@ def binary_cross_entropy_with_logits(x: torch.Tensor, y: torch.Tensor) -> torch.
     return F.binary_cross_entropy_with_logits(x, y.float())
 
 
-class ClassificationTask(Task):
+class ClassificationMixin:
+    def _build(
+        self,
+        num_classes: Optional[int] = None,
+        loss_fn: Optional[Callable] = None,
+        metrics: Union[torchmetrics.Metric, Mapping, Sequence, None] = None,
+        multi_label: bool = False,
+    ):
+        if metrics is None:
+            metrics = torchmetrics.F1(num_classes) if (multi_label and num_classes) else torchmetrics.Accuracy()
+
+        if loss_fn is None:
+            loss_fn = binary_cross_entropy_with_logits if multi_label else F.cross_entropy
+
+        return metrics, loss_fn
+
+    def to_metrics_format(self, x: torch.Tensor) -> torch.Tensor:
+        if getattr(self.hparams, "multi_label", False):
+            return torch.sigmoid(x)
+        return torch.softmax(x, dim=1)
+
+
+class ClassificationTask(Task, ClassificationMixin):
     def __init__(
         self,
         *args,
@@ -48,11 +71,9 @@ def __init__(
         serializer: Optional[Union[Serializer, Mapping[str, Serializer]]] = None,
         **kwargs,
     ) -> None:
-        if metrics is None:
-            metrics = torchmetrics.F1(num_classes) if (multi_label and num_classes) else torchmetrics.Accuracy()
 
-        if loss_fn is None:
-            loss_fn = binary_cross_entropy_with_logits if multi_label else F.cross_entropy
+        metrics, loss_fn = ClassificationMixin._build(self, num_classes, loss_fn, metrics, multi_label)
+
         super().__init__(
             *args,
             loss_fn=loss_fn,
@@ -61,11 +82,28 @@ def __init__(
             **kwargs,
         )
 
-    def to_metrics_format(self, x: torch.Tensor) -> torch.Tensor:
-        if getattr(self.hparams, "multi_label", False):
-            return torch.sigmoid(x)
-        # we'll assume that the data always comes as `(B, C, ...)`
-        return torch.softmax(x, dim=1)
+
+class ClassificationAdapterTask(AdapterTask, ClassificationMixin):
+    def __init__(
+        self,
+        *args,
+        num_classes: Optional[int] = None,
+        loss_fn: Optional[Callable] = None,
+        metrics: Union[torchmetrics.Metric, Mapping, Sequence, None] = None,
+        multi_label: bool = False,
+        serializer: Optional[Union[Serializer, Mapping[str, Serializer]]] = None,
+        **kwargs,
+    ) -> None:
+
+        metrics, loss_fn = ClassificationMixin._build(self, num_classes, loss_fn, metrics, multi_label)
+
+        super().__init__(
+            *args,
+            loss_fn=loss_fn,
+            metrics=metrics,
+            serializer=serializer or Classes(multi_label=multi_label),
+            **kwargs,
+        )
 
 
 class ClassificationSerializer(Serializer):

flash/core/data/data_module.py

@@ -302,6 +302,7 @@ def _train_dataloader(self) -> DataLoader:
         if isinstance(getattr(self, "trainer", None), pl.Trainer):
             return self.trainer.lightning_module.process_train_dataset(
                 train_ds,
+                trainer=self.trainer,
                 batch_size=self.batch_size,
                 num_workers=self.num_workers,
                 pin_memory=pin_memory,
@@ -330,6 +331,7 @@ def _val_dataloader(self) -> DataLoader:
         if isinstance(getattr(self, "trainer", None), pl.Trainer):
             return self.trainer.lightning_module.process_val_dataset(
                 val_ds,
+                trainer=self.trainer,
                 batch_size=self.batch_size,
                 num_workers=self.num_workers,
                 pin_memory=pin_memory,
@@ -352,6 +354,7 @@ def _test_dataloader(self) -> DataLoader:
         if isinstance(getattr(self, "trainer", None), pl.Trainer):
             return self.trainer.lightning_module.process_test_dataset(
                 test_ds,
+                trainer=self.trainer,
                 batch_size=self.batch_size,
                 num_workers=self.num_workers,
                 pin_memory=pin_memory,

flash/core/data/data_pipeline.py

@@ -534,7 +534,7 @@ def _detach_preprocessing_from_model(self, model: "Task", stage: Optional[Runnin
                     if isinstance(dl_args["collate_fn"], _Preprocessor):
                         dl_args["collate_fn"] = dl_args["collate_fn"]._original_collate_fn
 
-                        if isinstance(dl_args["dataset"], IterableAutoDataset):
+                        if isinstance(dl_args["dataset"], (IterableAutoDataset, IterableDataset)):
                             del dl_args["sampler"]
 
                         del dl_args["batch_sampler"]

flash/core/data/data_source.py

@@ -464,6 +464,9 @@ def load_data(
 
             data = make_dataset(data, class_to_idx, extensions=self.extensions)
             return [{DefaultDataKeys.INPUT: input, DefaultDataKeys.TARGET: target} for input, target in data]
+        elif dataset is not None:
+            dataset.num_classes = len(np.unique(data[1]))
+
         return list(
             filter(
                 lambda sample: has_file_allowed_extension(sample[DefaultDataKeys.INPUT], self.extensions),

flash/core/data/process.py

@@ -342,17 +342,22 @@ def default_transforms() -> Optional[Dict[str, Callable]]:
         """
         return None
 
+    def _apply_sample_transform(self, sample: Any) -> Any:
+        if isinstance(sample, list):
+            return [self.current_transform(s) for s in sample]
+        return self.current_transform(sample)
+
     def pre_tensor_transform(self, sample: Any) -> Any:
         """Transforms to apply on a single object."""
-        return self.current_transform(sample)
+        return self._apply_sample_transform(sample)
 
     def to_tensor_transform(self, sample: Any) -> Tensor:
         """Transforms to convert single object to a tensor."""
-        return self.current_transform(sample)
+        return self._apply_sample_transform(sample)
 
     def post_tensor_transform(self, sample: Tensor) -> Tensor:
         """Transforms to apply on a tensor."""
-        return self.current_transform(sample)
+        return self._apply_sample_transform(sample)
 
     def per_batch_transform(self, batch: Any) -> Any:
         """Transforms to apply to a whole batch (if possible use this for efficiency).

flash/core/data/transforms.py

@@ -111,6 +111,8 @@ def kornia_collate(samples: Sequence[Dict[str, Any]]) -> Dict[str, Any]:
     This function removes that dimension and then
     applies ``torch.utils.data._utils.collate.default_collate``.
     """
+    if len(samples) == 1 and isinstance(samples[0], list):
+        samples = samples[0]
     for sample in samples:
         for key in sample.keys():
             if torch.is_tensor(sample[key]) and sample[key].ndim == 4:

flash/core/integrations/icevision/adapter.py

@@ -16,6 +16,7 @@
 
 from torch.utils.data import DataLoader, Sampler
 
+import flash
 from flash.core.adapter import Adapter
 from flash.core.data.auto_dataset import BaseAutoDataset
 from flash.core.data.data_source import DefaultDataKeys
@@ -91,6 +92,7 @@ def _collate_fn(collate_fn, samples, metadata: Optional[List[Dict[str, Any]]] =
     def process_train_dataset(
         self,
         dataset: BaseAutoDataset,
+        trainer: "flash.Trainer",
         batch_size: int,
         num_workers: int,
         pin_memory: bool,
@@ -114,6 +116,7 @@ def process_train_dataset(
     def process_val_dataset(
         self,
         dataset: BaseAutoDataset,
+        trainer: "flash.Trainer",
         batch_size: int,
         num_workers: int,
         pin_memory: bool,
@@ -137,6 +140,7 @@ def process_val_dataset(
     def process_test_dataset(
         self,
         dataset: BaseAutoDataset,
+        trainer: "flash.Trainer",
         batch_size: int,
         num_workers: int,
         pin_memory: bool,

flash/core/model.py

@@ -68,7 +68,7 @@ def __init__(self):
         self._children = []
 
         # TODO: create enum values to define what are the exact states
-        self._data_pipeline_state: Optional[DataPipelineState] = None
+        self._data_pipeline_state: DataPipelineState = DataPipelineState()
 
         # model own internal state shared with the data pipeline.
         self._state: Dict[Type[ProcessState], ProcessState] = {}
@@ -133,6 +133,7 @@ def _process_dataset(
     def process_train_dataset(
         self,
         dataset: BaseAutoDataset,
+        trainer: "flash.Trainer",
         batch_size: int,
         num_workers: int,
         pin_memory: bool,
@@ -155,6 +156,7 @@ def process_train_dataset(
     def process_val_dataset(
         self,
         dataset: BaseAutoDataset,
+        trainer: "flash.Trainer",
         batch_size: int,
         num_workers: int,
         pin_memory: bool,
@@ -177,6 +179,7 @@ def process_val_dataset(
     def process_test_dataset(
         self,
         dataset: BaseAutoDataset,
+        trainer: "flash.Trainer",
         batch_size: int,
         num_workers: int,
         pin_memory: bool,
@@ -422,6 +425,7 @@ def predict(
         else:
             x = self.transfer_batch_to_device(x, self.device)
         x = data_pipeline.device_preprocessor(running_stage)(x)
+        x = x[0] if isinstance(x, list) else x
         predictions = self.predict_step(x, 0)  # batch_idx is always 0 when running with `model.predict`
         predictions = data_pipeline.postprocessor(running_stage)(predictions)
         return predictions

flash/core/registry.py

@@ -88,7 +88,7 @@ def get(
         """
         matches = [e for e in self.functions if key == e["name"]]
         if not matches:
-            raise KeyError(f"Key: {key} is not in {type(self).__name__}")
+            raise KeyError(f"Key: {key} is not in {type(self).__name__}. Available keys: {self.available_keys()}")
 
         if metadata:
             matches = [m for m in matches if metadata.items() <= m["metadata"].items()]

flash/core/utilities/imports.py

@@ -98,6 +98,7 @@ def _compare_version(package: str, op, version) -> bool:
 _DATASETS_AVAILABLE = _module_available("datasets")
 _ICEVISION_AVAILABLE = _module_available("icevision")
 _ICEDATA_AVAILABLE = _module_available("icedata")
+_LEARN2LEARN_AVAILABLE = _module_available("learn2learn") and _compare_version("learn2learn", operator.ge, "0.1.6")
 _TORCH_ORT_AVAILABLE = _module_available("torch_ort")
 _VISSL_AVAILABLE = _module_available("vissl") and _module_available("classy_vision")
 

flash/core/utilities/providers.py

@@ -39,6 +39,7 @@ def __str__(self):
 _SEGMENTATION_MODELS = Provider(
     "qubvel/segmentation_models.pytorch", "https://github.com/qubvel/segmentation_models.pytorch"
 )
+_LEARN2LEARN = Provider("learnables/learn2learn", "https://github.com/learnables/learn2learn")
 _PYSTICHE = Provider("pystiche/pystiche", "https://github.com/pystiche/pystiche")
 _HUGGINGFACE = Provider("Hugging Face/transformers", "https://github.com/huggingface/transformers")
 _FAIRSEQ = Provider("PyTorch/fairseq", "https://github.com/pytorch/fairseq")