MetaBatch is a micro-framework for meta-learning in PyTorch. It provides convenient Taskset and
TaskLoader classes for batch-aware online task creation for meta-learning.
Training meta-learning models efficiently can be a challenge, especially when it comes to creating random tasks of a consistent shape in one batch. The task creation process can be time-consuming and typically requires all tasks in the batch to have the same amount of context and target points. This can be a bottleneck during training:
# Sample code for creating a batch of tasks with traditional approach
class MyTaskDataset(Dataset):
...
def __getitem__(self, idx):
task = self.task_data[idx]
return task
class Model(Module):
...
def forward(self, tasks):
ctx_batch = tasks['context']
tgt_batch = tasks['target']
...
# create dataset
task_data = [{'images': [...], 'label': 'dog'},
{'images': [...], 'label': 'cat'}, ...]
dataset = MyTaskDataset(task_data)
dataloader = DataLoader(dataset, batch_size=16, workers=8)
for batch in dataloader:
...
# Construct batch of random tasks in the training loop (bottleneck!)
n_context = random.randint(low=1, high=5)
n_target = random.randint(low=1, high=10)
tasks = {'context': [], 'target': []}
for task in batch:
context_images = sample_n_images(task['images'], n_context)
target_images = sample_n_images(task['images'], n_target)
tasks['context'].append(context_images)
tasks['target'].append(target_images)
model(tasks)
...Wouldn't it be better to offload the task creation to the dataloader, so that it can be done in
parallel on multiple cores?
With MetaBatch, we simplify the process by allowing you to do just that.
We provide a TaskSet wrapper, where you can implement the __gettask__(self, index, n_context, n_target)__ method instead of PyTorch's __getitem(self, index)__. Our TaskLoader and
custom sampler take care of synchronizing n_context and n_target for each batch element
dispatched to all workers. With MetaBatch, the training bottleneck can be removed from the
above example:
# Sample code for creating a batch of tasks with MetaBatch
from metabatch import TaskSet, TaskLoader
class MyTaskSet(TaskSet):
...
def __gettask__(self, idx, n_context, n_target):
data = self.task_data[idx]
context_images = sample_n_images(data['images'], n_context)
target_images = sample_n_images(data['images'], n_target)
return {
'context': context_images
'target': target_images
}
class Model(Module):
...
def forward(self, tasks):
ctx_batch = tasks['context']
tgt_batch = tasks['target']
...
# create dataset
task_data = [{'images': [...], 'label': 'dog'},
{'images': [...], 'label': 'cat'}, ...]
dataset = MyTaskSet(task_data, min_pts=1, max_ctx_pts=5, max_tgt_pts=10)
dataloader = TaskLoader(dataset, batch_size=16, workers=8)
for batch in dataloader:
...
# Simply access the batch of constructed tasks (no bottleneck!)
model(batch)
...Install it: pip install metabatch
Requirements:
pytorch
Look at the example above for an idea or how to use TaskLoader with TaskSet, or go through the
examples in examples/ (TODO).
- MetaBatch allows for efficient task creation and batching during training, resulting in more task variations since you are no longer limited to precomputed tasks.
- Reduces boilerplate needed to precompute and load tasks.
MetaBatch is a micro-framework for meta-learning in PyTorch that provides convenient tools for (potentially faster) meta-training. It simplifies the task creation process and allows for efficient batching, making it a useful tool for researchers and engineers working on meta-learning projects.
TODO: benchmark MAML and CNP examples with typical implementation and other repos.
MetaBatch is released under the MIT License. See the LICENSE file for more information.