Code for the ICML 2021 paper "Bridging Multi-Task Learning and Meta-Learning: Towards Efficient Training and Effective Adaptation", Haoxiang Wang, Han Zhao, Bo Li.

Overview

Bridging Multi-Task Learning and Meta-Learning

Code for the ICML 2021 paper "Bridging Multi-Task Learning and Meta-Learning: Towards Efficient Training and Effective Adaptation" by Haoxiang Wang, Han Zhao, and Bo Li from UIUC. [arXiv:2106.09017]

This repo contains our efficient implementation of multi-task learning (MTL) for few-shot learning benchmarks. Our empirical studies show that our efficient MTL implementation can match the state-of-the-art gradient-based meta-learning algorithms, while enjoying an order of magnitude less training time.

If you find this repo useful for your research, please consider citing our paper

@inproceedings{wang2021bridging,
  title={Bridging Multi-Task Learning and Meta-Learning: Towards Efficient Training and Effective Adaptation},
  author={Wang, Haoxiang and Zhao, Han and Li, Bo},
  booktitle={International Conference on Machine Learning},
  year={2021},
  organization={PMLR}
}

Installation

This repo was tested with Ubuntu 18.04, Python 3.7 & 3.8 (Anaconda version), Pytorch 1.8 with CUDA 10 & 11.

Required Packages

  • Common Packages (covered by Anaconda): numpy, scikit-learn, tqdm, etc.
  • Pytorch packages: pytorch, torchvision, pytorch-lightning (a PyTorch wrapper), pytorch-optimizer (a collection of optimizers for Pytorch)
  • Other packages: learn2learn (a meta-learning codebase), WandB (a smart logger to visualize and track machine learning experiments)

Installation Tutorial:

  1. Install the Conda package manager with Anaconda following this webpage.

  2. In the Conda environment, install PyTorch and TorchVision following the official tutorial.

  3. In the Conda environment, install all other packages using the command line

    pip install pytorch-lightning==1.0.2 torch_optimizer wandb 
    
    • Currently, learn2learn is not compatible with the latest version of pytorch-lightning.
  4. Install the latest version of learn2learn by downloading its GitHub repo to a local folder

    git clone https://github.com/learnables/learn2learn.git && cd learn2learn && pip install -e .
    

Code

This repo is mainly built upon the learn2learn package (especially its pytorch-lightning version).

  • train.py: The script to train multi-task learning (and other meta-learning algorithms) on few-shot image classification benchmarks. This code is built upon this example provided by learn2learn.
  • lightning_episodic_module.py: This contains a base class LightningEpisodicModule for meta-learning. Notice that this class is built upon this implementation in learn2learn, and we add more choices of test methods, optimizers, and learning rate schedulers to this class.
  • lightning_mtl.py: This contains a class LightningMTL, which is our implementation of multi-task learning (MTL). Notice that this is a child class of LightningEpisodicModule, and it only includes the training code. Please refer to lightning_episodic_module.py for the optimizer setup and the test/validation evaluation code (e.g., fine-tuning the last layer by L2-regularized logistic regression)
  • models/: This folder is adopted from the repo RFS by Yue Wang, and it contains the implementation of multiple variants of ResNet. For a fair comparison with previous few-shot learning works, we conducted all experiments with the ResNet-12 architecture. However, one can try other architectures (e.g., ResNet-50 or WideResNet-28) to obtain possibly better performance.
  • datasets/: This is adapted from this learn2learn folder, which contains data loaders for few-shot image classification datasets. We only made a slight modification to two data loaders (mini-ImageNet and tiered-ImageNet), so we only present them here. Other data-loaders can be directly loaded from the learn2learn package.
  • benchmarks/: This is adapted from this learn2learn folder, which provides wrappers for data loaders (e.g., few-shot task construction, data augmentation). We implemented the data augmentation introduced in MetaOptNet (named as Lee19 in our code), and added some functions to accommodate the training of multi-task learning.

Logging and Checkpointing

  • Logging: We use the WandB logger, which can automatically upload all training and evaluation logs to the cloud. You need to register a free WandB account first and log in to it on your workstation/server following this tutorial. You can also turn off the sync function to make the logger purely offline. Notice that the logger we used is wrapped by Pytorch-Lightning (see this doc), and you can easily switch it to other loggers supported by Pytorch-Lighting (e.g., TensorBoard, MLflow, Comet, Neptune), since the logging functions we use in this repo is universal to all Pytorch-Lightning loggers. See this tutorial for details.
  • Checkpointing: We adopt the checkpointing functions provided by Pytorch-Lightning, which can automatically save checkpoints at the end of every epoch. Our implementation only saves the best epoch checkpoint (i.e., the epoch with the highest validation accuracy) and the last epoch checkpoint, so the disk usage of checkpoints is constant. Pytorch-Lightning can also help you load checkpoints easily.

Args and Flags

Below, we introduce multiple arguments and flags you might want to modify.

Notice that the argparser does not only parse the arguments in train.py. In addition, it also parses arguments in 'lightning_episodic_module.py' and lightning_mtl.py. Furthermore, some arguments are embedded in pytorch-lightning (e.g., --max_epochs).

  • --dataset: Currently, we support four datasets, ['mini-imagenet', 'tiered-imagenet', 'cifarfs', 'fc100']
  • --algorithm: We only provide the implementation of multi-task learning in this repo. However, you can also try other meta-learning algorithms provided by this learn2learn folder.
    • Note: in the empirical comparison of our paper, we take the officially reported performance in the original papers of these meta-learning methods (e.g., MAML, MetaOptNet), instead of running this script with the learn2learn implementation.
  • --root: The root directory that saves datasets. The default is ~/data.
  • --log_dir: The directory that saves logs and checkpoints. The defult is ~/wandb_logs
  • --gpu: The index of the GPU you want to use. Currently, our code only supports single-GPU training.
  • --norm_train_features: Normalize features (i.e., last hidden layer outputs) during training. We find this is useful for all datasets except for FC100.
  • --test_method: The default is l2, which is the L2-regularized logistic regression we use for the fine-tuning adaptation during the test phase.
  • --optim: The optimizer to use (default is 'radam', i.e., RAdam). The choice list is ['sgd', 'sgdp', 'sgdw', 'adam', 'adamw', 'adamax', 'radam', 'adabound', 'adamp']. These optimizers are adopted from PyTorch and pytorch-optimizer.
  • --scheduler: The learning rate scheduler to use (default is 'plateau', i.e., ReduceLROnPlateau). The other choice is 'step', which is the StepLR scheduler.
  • --meta_batch_size: The number of training tasks in a batch during (meta-)training.
  • --train_queries:
  • --test_shots: The number of shots during the test.
    • We only consider 5-way classification in our experiments. If you want to vary it, you can change the flags --train_ways and --test_ways.
    • Our implementation of multi-task learning does not care about the number of shots during training, since it does not split a task into query and support data. Thus --train_shots is noneffective for multi-task learning. But it is useful to other meta-learning algorithms.
  • --test_queries: The number of query samples per class in each test task (default is 30).
    • In our experiments, we adopt test_queries = 50. But this has a large GPU memory consumption, e.g., it will exceed 11GB (the limit of 1080ti and 2080ti) when running on mini-ImageNet. Thus, we set it as 30 here instead. In principle, larger test_queries leads to a more accurate estimation of the test accuracy. However, if you cannot use a large value of test_queries, you can also enlarge the following final_test_epoch_length to obtain a more accurate test accuracy estimation.
  • --final_test_epoch_length: The number of tasks for test evaluation at the end of training. Default is 2000.
  • --no_log: Running the script with this flag will turn off the logging and checkpointing. Notice that in this mode, pytorch-lightning will automatically write some checkpoint files to ./checkpoints/.
  • --max_epochs: the number of maximum training epochs.

Running

To replicate empirical results of multi-task learning shown in the paper, please run the following commands with your preferred root directory to save datasets and the GPU to use. The code can automatically download datasets to the root directory you specified.

In the following demo, we consider the root directory as /data and the GPU index as 0.

mini-ImageNet

  • 5-shot

    python train.py --gpu 0 --root ~/data --norm_train_features --meta_batch_size 2 --test_shots 5 --algorithm mtl --dataset mini-imagenet
    
  • 1-shot

    python train.py --gpu 0 --root ~/data --norm_train_features  --test_shots 1 --algorithm mtl --dataset mini-imagenet
    

tiered-ImageNet

  • 5-shot

    python train.py --gpu 0 --root ~/data --norm_train_features --train_queries 35 --meta_batch_size 1 --test_shots 5 --algorithm mtl --dataset tiered-imagenet
    
    • For multi-task learning, there is no query-support split during training (i.e., one can think that we "merge" query and support data for multi-task learning). Then, given the default train_queries = 5, setting --train_queries 35 makes the number of samples per class in each training task to be 40.
  • 1-shot

    python train.py --gpu 0 --root ~/data --norm_train_features --train_queries 35 --meta_batch_size 2 --test_shots 1 --algorithm mtl --dataset tiered-imagenet
    

CIFAR-FS

  • 5-shot

    python train.py --gpu 0 --root ~/data --norm_train_features --test_shots 5 --algorithm mtl --dataset cifarfs
    
  • 1-shot

    python train.py --gpu 0 --root ~/data --norm_train_features --train_queries 15 --meta_batch_size 2 --test_shots 1 --algorithm mtl --dataset cifarfs
    

FC100

  • 5-shot

    python train.py --gpu 0 --root ~/data --test_shots 5 --algorithm mtl --dataset fc100
    
  • 1-shot

    python train.py --gpu 0 --root ~/data --train_queries 4 --train_shots 1 --meta_batch_size 1 --test_shots 1 --algorithm mtl --dataset fc100
    
    • Setting --train_queries 4 --train_shots 1 reduces the number of samples per class in each training task to be 4+1=5 for multi-task learning, since it does not have the query-support splitting.

Results

Here is an empirical comparison of our implementation of multi-task learning against MetaOptNet, a state-of-the-art gradient-based meta-learning algorithm. For more results, please refer to our paper.

Benchmarking

Efficiency comparison

We benchmarked the training cost of Multi-Task Learning and MetaOptNet on AWS-EC2 P3 instances. See our papers for more details.

  • mini-ImageNet (5-way 5-shot)
Test Accuracy GPU Hours
MetaOptNet 78.63% 85.6 hrs
Multi-Task Learning 77.72% 3.7 hrs
  • tiered-ImageNet (5-way 5-shot)

    efficiency-tiered-imagenet
    • We vary the meta_batch_size and the number of epochs and finally obtain this plot. Each dot of MTL represents the result of one configuration. With smaller meta_batch_size and the number of epochs, MTL can be 11x faster than MetaOptNet while achieving the same performance (81.55% vs. 81.56%).

Contact Information

Please contact Haoxiang Wang ([email protected]) for any questions regarding this repo.

If you have questions regarding the learn2learn codebase, please consider joining the learn2learn slack to ask the authors and contributors of learn2learn directly.

Acknowledgments

In this repo, we adopt some code from the following codebases, and we sincerely thank their authors:

  • learn2learn: This repo is built on the learn2learn codebase (more precisely, the pytorch-lightning version of learn2learn), and the implemented multi-task learning code will be pushed to learn2learn in the future. Stay tuned!
  • RFS: We adopt the implementation of neural net architectures (e.g., CNNs, ResNets) from the RFS codebase (i.e., we copied their models/ to the './models' in this repo).
Owner
AI Secure
UIUC Secure Learning Lab
AI Secure
Exploring Image Deblurring via Blur Kernel Space (CVPR'21)

Exploring Image Deblurring via Encoded Blur Kernel Space About the project We introduce a method to encode the blur operators of an arbitrary dataset

VinAI Research 118 Dec 19, 2022
A 1.3B text-to-image generation model trained on 14 million image-text pairs

minDALL-E on Conceptual Captions minDALL-E, named after minGPT, is a 1.3B text-to-image generation model trained on 14 million image-text pairs for no

Kakao Brain 604 Dec 14, 2022
This repo holds codes of the ICCV21 paper: Visual Alignment Constraint for Continuous Sign Language Recognition.

VAC_CSLR This repo holds codes of the paper: Visual Alignment Constraint for Continuous Sign Language Recognition.(ICCV 2021) [paper] Prerequisites Th

Yuecong Min 64 Dec 19, 2022
Official Implementation of Neural Splines

Neural Splines: Fitting 3D Surfaces with Inifinitely-Wide Neural Networks This repository contains the official implementation of the CVPR 2021 (Oral)

Francis Williams 56 Nov 29, 2022
IsoGCN code for ICLR2021

IsoGCN The official implementation of IsoGCN, presented in the ICLR2021 paper Isometric Transformation Invariant and Equivariant Graph Convolutional N

horiem 39 Nov 25, 2022
A framework for analyzing computer vision models with simulated data

3DB: A framework for analyzing computer vision models with simulated data Paper Quickstart guide Blog post Installation Follow instructions on: https:

3DB 112 Jan 01, 2023
Translation-equivariant Image Quantizer for Bi-directional Image-Text Generation

Translation-equivariant Image Quantizer for Bi-directional Image-Text Generation Woncheol Shin1, Gyubok Lee1, Jiyoung Lee1, Joonseok Lee2,3, Edward Ch

Woncheol Shin 7 Sep 26, 2022
Implementation of E(n)-Transformer, which extends the ideas of Welling's E(n)-Equivariant Graph Neural Network to attention

E(n)-Equivariant Transformer (wip) Implementation of E(n)-Equivariant Transformer, which extends the ideas from Welling's E(n)-Equivariant G

Phil Wang 132 Jan 02, 2023
The full training script for Enformer (Tensorflow Sonnet) on TPU clusters

Enformer TPU training script (wip) The full training script for Enformer (Tensorflow Sonnet) on TPU clusters, in an effort to migrate the model to pyt

Phil Wang 10 Oct 19, 2022
[CVPR 2022 Oral] Rethinking Minimal Sufficient Representation in Contrastive Learning

Rethinking Minimal Sufficient Representation in Contrastive Learning PyTorch implementation of Rethinking Minimal Sufficient Representation in Contras

36 Nov 23, 2022
Unrolled Generative Adversarial Networks

Unrolled Generative Adversarial Networks Luke Metz, Ben Poole, David Pfau, Jascha Sohl-Dickstein arxiv:1611.02163 This repo contains an example notebo

Ben Poole 292 Dec 06, 2022
EmoTag helps you train emotion detection model for Chinese audios

emoTag emoTag helps you train emotion detection model for Chinese audios. Environment pip install -r requirement.txt Data We used Emotional Speech Dat

_zza 4 Sep 07, 2022
PointNetVLAD: Deep Point Cloud Based Retrieval for Large-Scale Place Recognition, CVPR 2018

PointNetVLAD: Deep Point Cloud Based Retrieval for Large-Scale Place Recognition PointNetVLAD: Deep Point Cloud Based Retrieval for Large-Scale Place

Mikaela Uy 294 Dec 12, 2022
Pixel-level Crack Detection From Images Of Levee Systems : A Comparative Study

PIXEL-LEVEL CRACK DETECTION FROM IMAGES OF LEVEE SYSTEMS : A COMPARATIVE STUDY G

Manisha Panta 2 Jul 23, 2022
HGCAE Pytorch implementation. CVPR2021 accepted.

Hyperbolic Graph Convolutional Auto-Encoders Accepted to CVPR2021 🎉 Official PyTorch code of Unsupervised Hyperbolic Representation Learning via Mess

Junho Cho 37 Nov 13, 2022
A simple baseline for 3d human pose estimation in PyTorch.

3d_pose_baseline_pytorch A PyTorch implementation of a simple baseline for 3d human pose estimation. You can check the original Tensorflow implementat

weigq 312 Jan 06, 2023
Repository for MDPGT

MD-PGT Repository for implementing and reproducing the results for the paper MDPGT: Momentum-based Decentralized Policy Gradient Tracking. Available E

Xian Yeow Lee 2 Dec 30, 2021
CS_Final_Metal_surface_detection - This is a final project for CoderSchool Machine Learning bootcamp on 29/12/2021.

CS_Final_Metal_surface_detection This is a final project for CoderSchool Machine Learning bootcamp on 29/12/2021. The project is based on the dataset

Cuong Vo 1 Dec 29, 2021
A Python Library for Graph Outlier Detection (Anomaly Detection)

PyGOD is a Python library for graph outlier detection (anomaly detection). This exciting yet challenging field has many key applications, e.g., detect

PyGOD Team 757 Jan 04, 2023
The code for 'Deep Residual Fourier Transformation for Single Image Deblurring'

Deep Residual Fourier Transformation for Single Image Deblurring Xintian Mao, Yiming Liu, Wei Shen, Qingli Li and Yan Wang code will be released soon

145 Dec 13, 2022