FFNet: MetaMixer-based Efficient Convolutional Mixer Design

Official PyTorch implementation of FFNet, from the following paper "FFNet: MetaMixer-based Efficient Convolutional Mixer Design".

Seokju Yun, Dongheon Lee, Youngmin Ro.

Figure: Overview of MetaMixer. (a) MetaMixer is derived by not specifying sub-operations within the query-key-value framework. We assert that the competence of Transformers primarily originates from MetaMixer, which we deem as the true backbone of Transformer. (b) To demonstrate this and propose a FFN-like efficient token mixer, we replace the inefficient sub-operations of self-attention with those from FFN while retaining MetaMixer structure. (c) Our MetaMixer-based pure ConvNets outperform domain-specialized competitors in various tasks, confirming the superiority of the MetaMixer framework.

Pre-trained Models

ImageNet-1K

Models trained on ImageNet-1K with knowledge distillation.

Training

Image Classification

Setup

conda create -n ffnet python=3.9
conda activate ffnet
conda install pytorch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0 cudatoolkit=11.3 -c pytorch
pip install -r requirements.txt

Dataset Preparation

Download the ImageNet-1K dataset and structure the data as follows:

/path/to/imagenet-1k/
  train/
    class1/
      img1.jpeg
    class2/
      img2.jpeg
  validation/
    class1/
      img3.jpeg
    class2/
      img4.jpeg

To train FFNet models, follow the respective command below:

# Without Distillation
python -m torch.distributed.launch --nproc_per_node=8 train.py \
/path/to/ImageNet/dataset --model ffnet_1 -b 128 --lr 1e-3 \
--native-amp --mixup 0.2 --output /path/to/save/results \
--input-size 3 256 256 --drop-path 0.1

# With Distillation
python -m torch.distributed.launch --nproc_per_node=8 train.py \
/path/to/ImageNet/dataset --model ffnet_1 -b 128 --lr 1e-3 \ 
--native-amp --output /path/to/save/results \
--input-size 3 256 256 --drop-path 0.02 \
--distillation-type "hard"

# Without Distillation
python -m torch.distributed.launch --nproc_per_node=8 train.py \
/path/to/ImageNet/dataset --model ffnet_2 -b 128 --lr 1e-3 \
--native-amp --mixup 0.2 --output /path/to/save/results \
--input-size 3 256 256 --drop-path 0.15

# With Distillation
python -m torch.distributed.launch --nproc_per_node=8 train.py \
/path/to/ImageNet/dataset --model ffnet_2 -b 128 --lr 1e-3 \ 
--native-amp --output /path/to/save/results \
--input-size 3 256 256 --drop-path 0.08 \
--distillation-type "hard"

# Without Distillation
python -m torch.distributed.launch --nproc_per_node=8 train.py \
/path/to/ImageNet/dataset --model ffnet_3 -b 128 --lr 1e-3 \
--native-amp --output /path/to/save/results \
--input-size 3 256 256 --drop-path 0.35

# With Distillation
python -m torch.distributed.launch --nproc_per_node=8 train.py \
/path/to/ImageNet/dataset --model ffnet_3 -b 128 --lr 1e-3 \ 
--native-amp --output /path/to/save/results \
--input-size 3 256 256 --drop-path 0.2 \
--distillation-type "hard"

We finetune models pre-trained on ImageNet-1K at a resolution of 384x384.

python -m torch.distributed.launch --nproc_per_node=8 train.py \
/path/to/ImageNet/dataset --model ffnet_3 -b 64 --lr 5e-6 \
--native-amp --output /path/to/save/results --input-size 3 384 384 \
--drop-path 0.4 --epochs 30 --warmup-epochs 0 --weight-decay 1e-8 \
--sched none --finetune --resume /path/to/checkpoint.pth

python -m torch.distributed.launch --nproc_per_node=8 train.py \
/path/to/ImageNet/dataset --model ffnet_4 -b 64 --lr 5e-6 \
--native-amp --output /path/to/save/results --input-size 3 384 384 \
--drop-path 0.5 --epochs 30 --warmup-epochs 0 --weight-decay 1e-8 \
--sched none --finetune --resume /path/to/checkpoint.pth

Evaluation

Run the following command to evaluate a pre-trained FFNet-2 on ImageNet-1K validation set with a single GPU:

python -m torch.distributed.launch --nproc_per_node=1 train.py \
/path/to/ImageNet/dataset --model ffnet_2 -b 128 \
--native-amp --input-size 3 256 256 --resume /path/to/checkpoint.pth --eval

Latency Measurement

The mobile latency reported in our work for iPhone 12 uses the deployment tool from XCode 14.

export the model to Core ML model

python export_model.py --variant ffnet_1 --output-dir /path/to/save/exported_model \
--checkpoint /path/to/pretrained_checkpoints/ffnet_1.pth.tar

MetaMixer's Task-Generality

For detailed documentation, please refer to these documents and codes as follows:

• Detection and instance segmentation
• 2D semantic segmentation
• Super-resolution
• Time series long-term forecasting
• 3D semantic segmentation

Bibtex

@article{yun2024metamixer,
  title={FFNet: MetaMixer-based Efficient Convolutional Mixer Design},
  author={Yun, Seokju and Lee, Dongheon and Ro, Youngmin},
  journal={arXiv preprint arXiv:2406.02021},
  year={2024}
}

License

This project is released under the Apache 2.0 license. Please see the LICENSE file for more information.

Acknowledgment

We sincerely appreciate pytorch-image-models, PyTorch, DeiT, RepVGG, and FastViT for their wonderful implementations.