Package: metabatch
IntroductionMetaBatch is a micro-framework for meta-learning in PyTorch. It provides convenientTasksetandTaskLoaderclasses forbatch-aware online task creation for meta-learning.Efficient batchingTraining 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 approachclassMyTaskDataset(Dataset):...def__getitem__(self,idx):task=self.task_data[idx]returntaskclassModel(Module):...defforward(self,tasks):ctx_batch=tasks['context']tgt_batch=tasks['target']...# create datasettask_data=[{'images':[...],'label':'dog'},{'images':[...],'label':'cat'},...]dataset=MyTaskDataset(task_data)dataloader=DataLoader(dataset,batch_size=16,workers=8)forbatchindataloader:...# 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':[]}fortaskinbatch: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)...MultiprocessingWouldn't it be better to offload the task creation to the dataloader, so that it can be done in
parallel on multiple cores?
WithMetaBatch, we simplify the process by allowing you to do just that.
We provide aTaskSetwrapper, where you can implement the__gettask__(self, index, n_context, n_target)__method instead of PyTorch's__getitem(self, index)__. OurTaskLoaderand
custom sampler take care of synchronizingn_contextandn_targetfor each batch element
dispatched to all workers. WithMetaBatch, the training bottleneck can be removed from the
above example:# Sample code for creating a batch of tasks with MetaBatchfrommetabatchimportTaskSet,TaskLoaderclassMyTaskSet(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}classModel(Module):...defforward(self,tasks):ctx_batch=tasks['context']tgt_batch=tasks['target']...# create datasettask_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)forbatchindataloader:...# Simply access the batch of constructed tasks (no bottleneck!)model(batch)...Installation & usageInstall it:pip install metabatchRequirements:pytorchLook at the example above for an idea or how to useTaskLoaderwithTaskSet, or go through the
examples inexamples/(TODO).AdvantagesMetaBatch 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.How much faster?TODO: benchmark MAML and CNP examples with typical implementation and other repos.LicenseMetaBatch is released under the MIT License. See the LICENSE file for more information.