[PoC] Add MetaLearning support through learn2learn

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
@@ -68,6 +69,37 @@ def to_metrics_format(self, x: torch.Tensor) -> torch.Tensor:
         return torch.softmax(x, dim=1)
 
 
+class ClassificationAdapterTask(AdapterTask):
+    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:
+        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
+        super().__init__(
+            *args,
+            loss_fn=loss_fn,
+            metrics=metrics,
+            serializer=serializer or Classes(multi_label=multi_label),
+            **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 ClassificationSerializer(Serializer):
     """A base class for classification serializers.
 

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

@@ -344,14 +344,20 @@ def default_transforms() -> Optional[Dict[str, Callable]]:
 
     def pre_tensor_transform(self, sample: Any) -> Any:
         """Transforms to apply on a single object."""
+        if isinstance(sample, list):
+            return [self.current_transform(s) for s in sample]
         return self.current_transform(sample)
 
     def to_tensor_transform(self, sample: Any) -> Tensor:
         """Transforms to convert single object to a tensor."""
+        if isinstance(sample, list):
+            return [self.current_transform(s) for s in sample]
         return self.current_transform(sample)
 
     def post_tensor_transform(self, sample: Tensor) -> Tensor:
         """Transforms to apply on a tensor."""
+        if isinstance(sample, list):
+            return [self.current_transform(s) for s in sample]
         return self.current_transform(sample)
 
     def per_batch_transform(self, batch: Any) -> Any:

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/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")
 _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")