Add TasksetSampler (compatible with non-classification tasks) (#255)

* Add l2l.utils.warn_once

* Replace TaskDataset by Taskset everywhere.

* Add TasksetSampler implementation.

* Update docs and changelog.

* Bump to 0.2.0.

Author: seba-1511

CHANGELOG.md

@@ -15,11 +15,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Add `l2l.nn.MetaModule` and `l2l.nn.ParameterTransform` for parameter-efficient finetuning.
 * Add `l2l.nn.freeze`and `l2l.nn.unfreeze`.
 * Add Adapters and LoRA examples.
+* Add TasksetSampler, compatible with PyTorch's Dataloaders.
 
 ### Changed
 
 * Documentation: uses `mkdocstrings` instead of `pydoc-markdown`.
 * Remove `text/news_topic_classification.py` example.
+* Rename TaskDataset to Taskset.
 
 ### Fixed
 

README.md

@@ -19,7 +19,7 @@ To learn more, see our whitepaper: [arXiv:2008.12284](https://arxiv.org/abs/2008
 
 **Overview**
 
-* [`learn2learn.data`](http://learn2learn.net/docs/learn2learn.data/): `TaskDataset` and transforms to create few-shot tasks from any PyTorch dataset.
+* [`learn2learn.data`](http://learn2learn.net/docs/learn2learn.data/): `Taskset` and transforms to create few-shot tasks from any PyTorch dataset.
 * [`learn2learn.vision`](http://learn2learn.net/docs/learn2learn.vision/): Models, datasets, and benchmarks for computer vision and few-shot learning.
 * [`learn2learn.gym`](http://learn2learn.net/docs/learn2learn.gym/): Environment and utilities for meta-reinforcement learning.
 * [`learn2learn.algorithms`](http://learn2learn.net/docs/learn2learn.algorithms/): High-level wrappers for existing meta-learning algorithms.
@@ -101,7 +101,7 @@ transforms = [  # Easy to define your own transform
     l2l.data.transforms.KShots(dataset, k=1),
     l2l.data.transforms.LoadData(dataset),
 ]
-taskset = TaskDataset(dataset, transforms, num_tasks=20000)
+taskset = Taskset(dataset, transforms, num_tasks=20000)
 for task in taskset:
     X, y = task
     # Meta-train on the task

docs/docs/learn2learn.data.md

@@ -21,7 +21,7 @@
         - __init__
         - __getitem__
 
-::: learn2learn.data.TaskDataset
+::: learn2learn.data.Taskset
     selection:
       members:
         - __init__

docs/tutorials/anil_tutorial/ANIL_tutorial.md

@@ -93,12 +93,12 @@ train_transforms = [
     RemapLabels(train_dataset),
     ConsecutiveLabels(train_dataset),
 ]
-train_tasks = l2l.data.TaskDataset(train_dataset,
+train_tasks = l2l.data.Taskset(train_dataset,
                                    task_transforms=train_transforms,
                                    num_tasks=20000)
 ~~~
 
-`l2l.data.TaskDataset` creates a set of tasks from the MetaDataset using a list of task transformations:
+`l2l.data.Taskset` creates a set of tasks from the MetaDataset using a list of task transformations:
 
 * `FusedNWaysKShots(dataset, n=ways, k=2*shots)`: efficient implementation to keep \(k\) data samples from \(n\) randomly sampled labels.
 

docs/tutorials/task_transform_tutorial/transform_tutorial.md

@@ -11,7 +11,7 @@ In this tutorial, we will explore in depth one of the core utilities [learn2lear
 *   We will first discuss the motivation behind generating tasks. *(Those familiar with meta-learning can skip this section.)*
 *   Next, we will have a high-level overview of the overall pipeline used for generating tasks using `learn2learn`.
 *   `MetaDataset` is used fast indexing, and accelerates the process of generating few-shot learning tasks. `UnionMetaDataset` and `FilteredMetaDataset` are extensions of `MetaDataset` that can further provide customised utility. `UnionMetaDataset` builds up on `MetaDataset` to construct a union of multiple input datasets, and `FilteredMetaDataset` takes in a `MetaDataset` and filters it to include only the required labels.
-*   `TaskDataset` is the core module that generates tasks from input dataset. Tasks are lazily sampled upon indexing or calling `.sample()` method.
+*   `Taskset` is the core module that generates tasks from input dataset. Tasks are lazily sampled upon indexing or calling `.sample()` method.
 *   Lastly, we study different `task transforms` defined in `learn2learn` that modifies the input data such that a customised `task` is generated. 
 
 ## Motivation for generating tasks
@@ -109,7 +109,7 @@ transforms = [
                     ]
 
 # 5. Generate set of tasks using the dataset, and transforms
-taskset = l2l.data.TaskDataset(dataset=omniglot, task_transforms=transforms, num_tasks=10) # Creates sets of tasks from the dataset 
+taskset = l2l.data.Taskset(dataset=omniglot, task_transforms=transforms, num_tasks=10) # Creates sets of tasks from the dataset 
 
 # Now sample a task from the taskset
 X, y = taskset.sample()
@@ -270,15 +270,15 @@ print('Original Labels:', len(toy_omniglot.labels))
 print('Filtered Labels:', len(filtered.labels))
 ~~~
 
-## TaskDataset - Core module
+## Taskset - Core module
 
 #### Introduction
 
 This is one of the core module of `learn2learn` that is used to generate a task from a given input dataset. It takes `dataset`, and list of `task transformations` as arguments. The task transformation basically define the kind of tasks that will be generated from the dataset. (For example, `KShots` transform limits the number of samples per class in a task to `K` samples per class.) 
 
 > If there are no task transforms, then the task consists of all the samples in the entire dataset.
 
-Another argument that `TaskDataset` takes as input is `num_tasks` *(an integer value)*. The value is set depending on how many tasks the user wants to generate. By default, it is kept as `-1`, meaning infinite number of tasks will be generated, and a new task is generated on sampling. In the former case, the descriptions of the task will be cached in a dictionary such that if a given task is called again, the description can be loaded instantly rather than generating it once again.
+Another argument that `Taskset` takes as input is `num_tasks` *(an integer value)*. The value is set depending on how many tasks the user wants to generate. By default, it is kept as `-1`, meaning infinite number of tasks will be generated, and a new task is generated on sampling. In the former case, the descriptions of the task will be cached in a dictionary such that if a given task is called again, the description can be loaded instantly rather than generating it once again.
 
 #### What is a task description?
 
@@ -627,7 +627,7 @@ toy_transforms = [
                 ConsecutiveLabels(omniglot), # Re-orders samples s.t. they are sorted in consecutive order 
                 RandomClassRotation(omniglot, [0, 90, 180, 270]) # Randomly rotate sample over x degrees (only for vision tasks)
                 ]
-toy_taskset = l2l.data.TaskDataset(omniglot, toy_transforms, num_tasks=20000)
+toy_taskset = l2l.data.Taskset(omniglot, toy_transforms, num_tasks=20000)
 try:
     print(len(toy_taskset.sample())) # Expected error as RemapLabels is used before LoadData
 except TypeError:
@@ -639,9 +639,9 @@ except TypeError:
 Traceback (most recent call last):
   File "<ipython-input-27-4c0558e6745b>", line 13, in <module>
     print(len(toy_taskset.sample())) # Expected error as RemapLabels is used before LoadData
-  File "learn2learn/data/task_dataset.pyx", line 158, in learn2learn.data.task_dataset.CythonTaskDataset.sample
-  File "learn2learn/data/task_dataset.pyx", line 173, in learn2learn.data.task_dataset.CythonTaskDataset.__getitem__
-  File "learn2learn/data/task_dataset.pyx", line 142, in learn2learn.data.task_dataset.CythonTaskDataset.get_task
+  File "learn2learn/data/task_dataset.pyx", line 158, in learn2learn.data.task_dataset.CythonTaskset.sample
+  File "learn2learn/data/task_dataset.pyx", line 173, in learn2learn.data.task_dataset.CythonTaskset.__getitem__
+  File "learn2learn/data/task_dataset.pyx", line 142, in learn2learn.data.task_dataset.CythonTaskset.get_task
   File "learn2learn/data/transforms.pyx", line 201, in learn2learn.data.transforms.RemapLabels.remap
 TypeError: 'int' object is not iterable
 ~~~

examples/vision/anil_fc100.py

@@ -96,29 +96,35 @@ def main(
         RemapLabels(train_dataset),
         ConsecutiveLabels(train_dataset),
     ]
-    train_tasks = l2l.data.TaskDataset(train_dataset,
-                                       task_transforms=train_transforms,
-                                       num_tasks=20000)
+    train_tasks = l2l.data.Taskset(
+        train_dataset,
+        task_transforms=train_transforms,
+        num_tasks=20000,
+    )
 
     valid_transforms = [
         FusedNWaysKShots(valid_dataset, n=ways, k=2*shots),
         LoadData(valid_dataset),
         ConsecutiveLabels(valid_dataset),
         RemapLabels(valid_dataset),
     ]
-    valid_tasks = l2l.data.TaskDataset(valid_dataset,
-                                       task_transforms=valid_transforms,
-                                       num_tasks=600)
+    valid_tasks = l2l.data.Taskset(
+        valid_dataset,
+        task_transforms=valid_transforms,
+        num_tasks=600,
+    )
 
     test_transforms = [
         FusedNWaysKShots(test_dataset, n=ways, k=2*shots),
         LoadData(test_dataset),
         RemapLabels(test_dataset),
         ConsecutiveLabels(test_dataset),
     ]
-    test_tasks = l2l.data.TaskDataset(test_dataset,
-                                      task_transforms=test_transforms,
-                                      num_tasks=600)
+    test_tasks = l2l.data.Taskset(
+        test_dataset,
+        task_transforms=test_transforms,
+        num_tasks=600,
+    )
 
 
     # Create model

examples/vision/meta_mnist.py

@@ -53,15 +53,17 @@ def main(lr=0.005, maml_lr=0.01, iterations=1000, ways=5, shots=1, tps=32, fas=5
                                              download=True,
                                              transform=transformations))
 
-    train_tasks = l2l.data.TaskDataset(mnist_train,
-                                       task_transforms=[
-                                            l2l.data.transforms.NWays(mnist_train, ways),
-                                            l2l.data.transforms.KShots(mnist_train, 2*shots),
-                                            l2l.data.transforms.LoadData(mnist_train),
-                                            l2l.data.transforms.RemapLabels(mnist_train),
-                                            l2l.data.transforms.ConsecutiveLabels(mnist_train),
-                                       ],
-                                       num_tasks=1000)
+    train_tasks = l2l.data.Taskset(
+        mnist_train,
+        task_transforms=[
+            l2l.data.transforms.NWays(mnist_train, ways),
+            l2l.data.transforms.KShots(mnist_train, 2*shots),
+            l2l.data.transforms.LoadData(mnist_train),
+            l2l.data.transforms.RemapLabels(mnist_train),
+            l2l.data.transforms.ConsecutiveLabels(mnist_train),
+        ],
+        num_tasks=1000,
+    )
 
     model = Net(ways)
     model.to(device)

examples/vision/protonet_miniimagenet.py

@@ -113,7 +113,7 @@ def fast_adapt(model, batch, ways, shot, query_num, metric=None, device=None):
         LoadData(train_dataset),
         RemapLabels(train_dataset),
     ]
-    train_tasks = l2l.data.TaskDataset(train_dataset, task_transforms=train_transforms)
+    train_tasks = l2l.data.Taskset(train_dataset, task_transforms=train_transforms)
     train_loader = DataLoader(train_tasks, pin_memory=True, shuffle=True)
 
     valid_dataset = l2l.data.MetaDataset(valid_dataset)
@@ -123,9 +123,11 @@ def fast_adapt(model, batch, ways, shot, query_num, metric=None, device=None):
         LoadData(valid_dataset),
         RemapLabels(valid_dataset),
     ]
-    valid_tasks = l2l.data.TaskDataset(valid_dataset,
-                                       task_transforms=valid_transforms,
-                                       num_tasks=200)
+    valid_tasks = l2l.data.Taskset(
+        valid_dataset,
+        task_transforms=valid_transforms,
+        num_tasks=200,
+    )
     valid_loader = DataLoader(valid_tasks, pin_memory=True, shuffle=True)
 
     test_dataset = l2l.data.MetaDataset(test_dataset)
@@ -135,9 +137,11 @@ def fast_adapt(model, batch, ways, shot, query_num, metric=None, device=None):
         LoadData(test_dataset),
         RemapLabels(test_dataset),
     ]
-    test_tasks = l2l.data.TaskDataset(test_dataset,
-                                      task_transforms=test_transforms,
-                                      num_tasks=2000)
+    test_tasks = l2l.data.Taskset(
+        test_dataset,
+        task_transforms=test_transforms,
+        num_tasks=2000,
+    )
     test_loader = DataLoader(test_tasks, pin_memory=True, shuffle=True)
 
     optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

learn2learn/_version.py

@@ -1 +1 @@
-__version__ = '0.1.7'
+__version__ = '0.2.0'

learn2learn/data/__init__.py

@@ -6,5 +6,5 @@
 
 from . import transforms
 from .meta_dataset import MetaDataset, UnionMetaDataset, FilteredMetaDataset
-from .task_dataset import TaskDataset, DataDescription
+from .task_dataset import TaskDataset, Taskset, DataDescription
 from .utils import OnDeviceDataset, partition_task, InfiniteIterator

learn2learn/data/meta_dataset.pyx

@@ -11,6 +11,7 @@ from collections import defaultdict
 import numpy as np
 import torch
 from torch.utils.data import Dataset
+import learn2learn as l2l
 
 
 class MetaDataset(Dataset):

learn2learn/data/samplers.py

@@ -0,0 +1,67 @@
+#!/usr/bin/env python
+
+import random
+import torch
+import learn2learn as l2l
+
+
+class TasksetSampler(torch.utils.data.Sampler):
+
+    def __init__(self, taskset, shuffle=True):
+        self.taskset = taskset
+        self.shuffle = shuffle
+
+    def description2indices(self, task_description):
+        return [dd.index for dd in task_description]
+
+    def __iter__(self):
+        if self.taskset.num_tasks == -1:  # loop infinitely
+            while True:
+                yield self.description2indices(
+                    self.taskset.sample_task_description()
+                )
+        else:  # loop over the range of tasks
+            task_indices = list(range(self.taskset.num_tasks))
+            if self.shuffle:
+                random.shuffle(task_indices)
+            for i in task_indices:
+                if i not in self.taskset.sampled_descriptions:
+                    self.taskset.sampled_descriptions[i] = self.taskset.sample_task_description()
+                yield self.description2indices(
+                    self.taskset.sampled_descriptions[i]
+                )
+
+
+if __name__ == "__main__":
+    NUM_TASKS = 10
+    NUM_DATA = 128
+    X_SHAPE = 16
+    Y_SHAPE = 10
+    EPSILON = 1e-6
+    SUBSET_SIZE = 5
+    WORKERS = 4
+    META_BSZ = 16
+    data = torch.randn(NUM_DATA, X_SHAPE)
+    labels = torch.randint(0, Y_SHAPE, (NUM_DATA, ))
+    dataset = torch.utils.data.TensorDataset(data, labels)
+    dataset = l2l.data.MetaDataset(dataset)
+    taskset = l2l.data.Taskset(
+        dataset,
+        task_transforms=[
+            l2l.data.transforms.FusedNWaysKShots(dataset, n=2, k=1),
+            l2l.data.transforms.LoadData(dataset),
+            l2l.data.transforms.RemapLabels(dataset),
+            l2l.data.transforms.ConsecutiveLabels(dataset),
+        ],
+        num_tasks=NUM_TASKS,
+    )
+
+    sampler = TasksetSampler(taskset)
+    dataloader = torch.utils.data.DataLoader(
+        dataset=dataset,
+        batch_sampler=sampler,
+    )
+    for task in dataloader:
+        print(task)
+
+    __import__('pdb').set_trace()

learn2learn/data/task_dataset.pyx

@@ -47,7 +47,7 @@ cdef class DataDescription:
         self.transforms = []
 
 
-class TaskDataset(CythonTaskDataset):
+class Taskset(CythonTaskDataset):
 
     """
     [[Source]](https://github.com/learnables/learn2learn/blob/master/learn2learn/data/task_dataset.py)
@@ -89,14 +89,24 @@ class TaskDataset(CythonTaskDataset):
     """
 
     def __init__(self, dataset, task_transforms=None, num_tasks=-1, task_collate=None):
-        super(TaskDataset, self).__init__(
+        super(Taskset, self).__init__(
             dataset=dataset,
             task_transforms=task_transforms,
             num_tasks=num_tasks,
             task_collate=task_collate,
         )
 
 
+class TaskDataset(Taskset):
+
+    def __init__(self, *args, **kwargs):
+        super(TaskDataset, self).__init__(*args, **kwargs)
+        l2l.utils.warn_once(
+            message='TaskDataset is deprecated, use Taskset instead.',
+            severity='deprecation',
+        )
+
+
 cdef class CythonTaskDataset:
 
     cdef public:

learn2learn/utils/__init__.py

@@ -4,6 +4,7 @@
 import torch
 import argparse
 import dataclasses
+import warnings
 
 
 def magic_box(x):
@@ -373,3 +374,27 @@ def __getattr__(self, *args, **kwargs):
 
     def __call__(self, *args, **kwargs):
         self.raise_import()
+
+
+class _SingleWarning(object):
+
+    def __init__(self):
+        self.warned_messages = []
+        self.warning_categories = {
+            'default': UserWarning,
+            'deprecation': DeprecationWarning,
+        }
+
+    def __call__(self, message, severity=None):
+        if message not in self.warned_messages:
+            if severity is None:
+                severity = 'default'
+            if severity == 'error':
+                raise RuntimeError(message)
+            elif isinstance(severity, str):
+                severity = self.warning_categories[severity]
+            warnings.warn(message, category=severity)
+            self.warned_messages.append(message)
+
+
+warn_once = _SingleWarning()