Adds OnDeviceDataset, including device option for benchmarks. (#253)

* Version bump.

* Fix merge conflict in _version.py

* Add scripts for supervised pretraining.

* Remove TODO.md

* Fix linting.

* Add spec to pretrained backbones.

* Add support for OnDevice in benchmarks.

* Update supervised pretraining example.

* Fix import in l2l.data.

* Fix lightning_anil_no_travis test.

* Fix scipy import.

* Remove supervised pretraining for OnDevice merge.

* Update changelog.

Author: seba-1511

CHANGELOG.md

@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * `l2l.nn.PrototypicalClassifier` and `l2l.nn.SVMClassifier`.
 * Add `l2l.vision.models.WRN28`.
 * Separate modules for `CNN4Backbone`, `ResNet12Backbone`, `WRN28Backbones` w/ pretrained weights.
+* Add `l2l.data.OnDeviceDataset` and implement `device` parameter for benchmarks.
 
 ### Changed
 

learn2learn/data/__init__.py

@@ -7,3 +7,4 @@
 from . import transforms
 from .meta_dataset import MetaDataset, UnionMetaDataset, FilteredMetaDataset
 from .task_dataset import TaskDataset, DataDescription
+from .utils import OnDeviceDataset, partition_task, InfiniteIterator

learn2learn/data/utils.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import torch
 import requests
 import tqdm
 
@@ -44,3 +45,98 @@ def save_response_content(response, destination):
         for chunk in response.iter_content(CHUNK_SIZE):
             if chunk:  # filter out keep-alive new chunks
                 f.write(chunk)
+
+
+class InfiniteIterator(object):
+
+    def __init__(self, dataloader):
+        self.dataloader = dataloader
+        self.iterator = iter(self.dataloader)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        while True:
+            try:
+                return next(self.iterator)
+            except StopIteration:
+                self.iterator = iter(self.dataloader)
+
+
+def partition_task(data, labels, shots=1, ways=None):
+    assert data.size(0) == labels.size(0)
+    unique_labels = labels.unique()
+    if ways is None:
+        ways = unique_labels.numel()
+    data_shape = data.shape[1:]
+    num_support = ways * shots
+    num_query = data.size(0) - num_support
+    assert num_query % ways == 0, 'Only query_shot == support_shot supported.'
+    query_shots = num_query // ways
+    support_data = torch.empty(
+        (num_support,) + data_shape,
+        device=data.device,
+        dtype=data.dtype,
+    )
+    support_labels = torch.empty(
+        num_support,
+        device=labels.device,
+        dtype=labels.dtype,
+    )
+    query_data = torch.empty(
+        (num_query, ) + data_shape,
+        device=data.device,
+        dtype=data.dtype,
+    )
+    query_labels = torch.empty(
+        num_query,
+        device=labels.device,
+        dtype=labels.dtype,
+    )
+    for i, label in enumerate(unique_labels):
+        support_start = i * shots
+        support_end = support_start + shots
+        query_start = i * query_shots
+        query_end = query_start + query_shots
+
+        # filter data
+        label_data = data[labels == label]  # TODO: fancy index makes a copy.
+        num_label_data = label_data.size(0)
+        assert num_label_data == shots + query_shots, \
+            'Only same number of query per label supported.'
+
+        # set value of labels
+        support_labels[support_start:support_end].fill_(label)
+        query_labels[query_start:query_end].fill_(label)
+
+        # set value of data
+        support_data[support_start:support_end].copy_(label_data[:shots])
+        query_data[query_start:query_end].copy_(label_data[shots:])
+
+    return (support_data, support_labels), (query_data, query_labels)
+
+
+class OnDeviceDataset(torch.utils.data.TensorDataset):
+
+    def __init__(self, dataset, device=None, transform=None):
+        data = []
+        labels = []
+        for x, y in dataset:
+            data.append(x.unsqueeze(0))
+            labels.append(y)
+        data = torch.cat(data, dim=0)
+        labels = torch.tensor(labels)
+        if device is not None:
+            data = data.to(device)
+            labels = labels.to(device)
+        super(OnDeviceDataset, self).__init__(data, labels)
+        self.transform = transform
+        if hasattr(dataset, '_bookkeeping_path'):
+            self._bookkeeping_path = dataset._bookkeeping_path
+
+    def __getitem__(self, index):
+        x, y = super(OnDeviceDataset, self).__getitem__(index)
+        if self.transform is not None:
+            x = self.transform(x)
+        return x, y

learn2learn/utils/__init__.py

@@ -2,6 +2,8 @@
 
 import copy
 import torch
+import argparse
+import dataclasses
 
 
 def magic_box(x):
@@ -313,6 +315,38 @@ def accuracy(preds, targets):
     return acc / preds.size(0)
 
 
+def flatten_config(args, prefix=None):
+    flat_args = dict()
+    if isinstance(args, argparse.Namespace):
+        args = vars(args)
+        return flatten_config(args)
+    elif not dataclasses.is_dataclass(args) and not isinstance(args, dict):
+        flat_args[prefix] = args
+        return flat_args
+    elif dataclasses.is_dataclass(args):
+        keys = dataclasses.fields(args)
+        def getvalue(x): getattr(args, x.name)
+    elif isinstance(args, dict):
+        keys = args.keys()
+        def getvalue(x): args[x]
+    else:
+        raise 'Unknown args'
+    for key in keys:
+        value = getvalue(key)
+        if prefix is None:
+            if isinstance(key, str):
+                prefix_child = key
+            elif isinstance(key, dataclasses.Field):
+                prefix_child = key.name
+            else:
+                raise 'Unknown key'
+        else:
+            prefix_child = prefix + '.' + key.name
+        flat_child = flatten_config(value, prefix=prefix_child)
+        flat_args.update(flat_child)
+    return flat_args
+
+
 class _ImportRaiser(object):
 
     def __init__(self, name, command):

learn2learn/vision/benchmarks/__init__.py

@@ -81,6 +81,7 @@ def get_tasksets(
     * **test_ways** (int, *optional*, default=5) - The number of classes per test tasks. Also used for validation tasks.
     * **test_samples** (int, *optional*, default=10) - The number of samples per test tasks. Also used for validation tasks.
     * **num_tasks** (int, *optional*, default=-1) - The number of tasks in each TaskDataset.
+    * **device** (torch.Device, *optional*, default=None) - If not None, tasksets are loaded as Tensors on `device`.
     * **root** (str, *optional*, default='~/data') - Where the data is stored.
 
     **Example**
@@ -96,15 +97,13 @@ def get_tasksets(
     """
     root = os.path.expanduser(root)
 
-    if device is not None:
-        raise NotImplementedError('Device other than None not implemented. (yet)')
-
     # Load task-specific data and transforms
     datasets, transforms = _TASKSETS[name](train_ways=train_ways,
                                            train_samples=train_samples,
                                            test_ways=test_ways,
                                            test_samples=test_samples,
                                            root=root,
+                                           device=device,
                                            **kwargs)
     train_dataset, validation_dataset, test_dataset = datasets
     train_transforms, validation_transforms, test_transforms = transforms

learn2learn/vision/benchmarks/cifarfs_benchmark.py

@@ -12,6 +12,7 @@ def cifarfs_tasksets(
     test_ways=5,
     test_samples=10,
     root='~/data',
+    device=None,
     **kwargs,
 ):
     """Tasksets for CIFAR-FS benchmarks."""
@@ -28,6 +29,19 @@ def cifarfs_tasksets(
                                                transform=data_transform,
                                                mode='test',
                                                download=True)
+    if device is not None:
+        train_dataset = l2l.data.OnDeviceDataset(
+            dataset=train_dataset,
+            device=device,
+        )
+        valid_dataset = l2l.data.OnDeviceDataset(
+            dataset=valid_dataset,
+            device=device,
+        )
+        test_dataset = l2l.data.OnDeviceDataset(
+            dataset=test_dataset,
+            device=device,
+        )
     train_dataset = l2l.data.MetaDataset(train_dataset)
     valid_dataset = l2l.data.MetaDataset(valid_dataset)
     test_dataset = l2l.data.MetaDataset(test_dataset)

learn2learn/vision/benchmarks/fc100_benchmark.py

@@ -12,6 +12,7 @@ def fc100_tasksets(
     test_ways=5,
     test_samples=10,
     root='~/data',
+    device=None,
     **kwargs,
 ):
     """Tasksets for FC100 benchmarks."""
@@ -28,6 +29,19 @@ def fc100_tasksets(
                                              transform=data_transform,
                                              mode='test',
                                              download=True)
+    if device is not None:
+        train_dataset = l2l.data.OnDeviceDataset(
+            dataset=train_dataset,
+            device=device,
+        )
+        valid_dataset = l2l.data.OnDeviceDataset(
+            dataset=valid_dataset,
+            device=device,
+        )
+        test_dataset = l2l.data.OnDeviceDataset(
+            dataset=test_dataset,
+            device=device,
+        )
     train_dataset = l2l.data.MetaDataset(train_dataset)
     valid_dataset = l2l.data.MetaDataset(valid_dataset)
     test_dataset = l2l.data.MetaDataset(test_dataset)

learn2learn/vision/benchmarks/mini_imagenet_benchmark.py

@@ -14,6 +14,7 @@ def mini_imagenet_tasksets(
     test_samples=10,
     root='~/data',
     data_augmentation=None,
+    device=None,
     **kwargs,
 ):
     """Tasksets for mini-ImageNet benchmarks."""
@@ -47,21 +48,38 @@ def mini_imagenet_tasksets(
     train_dataset = l2l.vision.datasets.MiniImagenet(
         root=root,
         mode='train',
-        transform=train_data_transforms,
         download=True,
     )
     valid_dataset = l2l.vision.datasets.MiniImagenet(
         root=root,
         mode='validation',
-        transform=test_data_transforms,
         download=True,
     )
     test_dataset = l2l.vision.datasets.MiniImagenet(
         root=root,
         mode='test',
-        transform=test_data_transforms,
         download=True,
     )
+    if device is None:
+        train_dataset.transform = train_data_transforms
+        valid_dataset.transform = test_data_transforms
+        test_dataset.transform = test_data_transforms
+    else:
+        train_dataset = l2l.data.OnDeviceDataset(
+            dataset=train_dataset,
+            transform=train_data_transforms,
+            device=device,
+        )
+        valid_dataset = l2l.data.OnDeviceDataset(
+            dataset=valid_dataset,
+            transform=test_data_transforms,
+            device=device,
+        )
+        test_dataset = l2l.data.OnDeviceDataset(
+            dataset=test_dataset,
+            transform=test_data_transforms,
+            device=device,
+        )
     train_dataset = l2l.data.MetaDataset(train_dataset)
     valid_dataset = l2l.data.MetaDataset(valid_dataset)
     test_dataset = l2l.data.MetaDataset(test_dataset)

learn2learn/vision/benchmarks/omniglot_benchmark.py

@@ -13,7 +13,8 @@ def omniglot_tasksets(
     test_ways,
     test_samples,
     root,
-    **kwargs
+    device=None,
+    **kwargs,
 ):
     """
     Benchmark definition for Omniglot.
@@ -28,18 +29,20 @@ def omniglot_tasksets(
         transform=data_transforms,
         download=True,
     )
+    if device is not None:
+        dataset = l2l.data.OnDeviceDataset(omniglot, device=device)
     dataset = l2l.data.MetaDataset(omniglot)
-    train_dataset = dataset
-    validation_datatset = dataset
-    test_dataset = dataset
 
     classes = list(range(1623))
     random.shuffle(classes)
+    train_dataset = l2l.data.FilteredMetaDataset(dataset, labels=classes[:1100])
+    validation_datatset = l2l.data.FilteredMetaDataset(dataset, labels=classes[1100:1200])
+    test_dataset = l2l.data.FilteredMetaDataset(dataset, labels=classes[1200:])
+
     train_transforms = [
         l2l.data.transforms.FusedNWaysKShots(dataset,
                                              n=train_ways,
-                                             k=train_samples,
-                                             filter_labels=classes[:1100]),
+                                             k=train_samples),
         l2l.data.transforms.LoadData(dataset),
         l2l.data.transforms.RemapLabels(dataset),
         l2l.data.transforms.ConsecutiveLabels(dataset),
@@ -48,8 +51,7 @@ def omniglot_tasksets(
     validation_transforms = [
         l2l.data.transforms.FusedNWaysKShots(dataset,
                                              n=test_ways,
-                                             k=test_samples,
-                                             filter_labels=classes[1100:1200]),
+                                             k=test_samples),
         l2l.data.transforms.LoadData(dataset),
         l2l.data.transforms.RemapLabels(dataset),
         l2l.data.transforms.ConsecutiveLabels(dataset),
@@ -58,8 +60,7 @@ def omniglot_tasksets(
     test_transforms = [
         l2l.data.transforms.FusedNWaysKShots(dataset,
                                              n=test_ways,
-                                             k=test_samples,
-                                             filter_labels=classes[1200:]),
+                                             k=test_samples),
         l2l.data.transforms.LoadData(dataset),
         l2l.data.transforms.RemapLabels(dataset),
         l2l.data.transforms.ConsecutiveLabels(dataset),