RepLKNet-pytorch (CVPR 2022)

This is the official PyTorch implementation of RepLKNet, from the following CVPR-2022 paper:

Scaling Up Your Kernels to 31x31: Revisiting Large Kernel Design in CNNs.

The paper is now released on arXiv: https://arxiv.org/abs/2203.06717.

Update: all the pretrained models, ImageNet-1K models, and Cityscapes/ADE20K/COCO models have been released.

Update: released a script to visualize the Effective Receptive Field (ERF). To get the ERF of your own model, you only need to add a few lines of code!

Update: released the training commands and more examples.

If you find the paper or this repository helpful, please consider citing

    @article{replknet,
    title={Scaling Up Your Kernels to 31x31: Revisiting Large Kernel Design in CNNs},
    author={Ding, Xiaohan and Zhang, Xiangyu and Zhou, Yizhuang and Han, Jungong and Ding, Guiguang and Sun, Jian},
    journal={arXiv preprint arXiv:2203.06717},
    year={2022}
    }

Other implementations

framework	link
MegEngine (official)	https://github.com/megvii-research/RepLKNet
PyTorch (official)	https://github.com/DingXiaoH/RepLKNet-pytorch
Tensorflow	https://github.com/shkarupa-alex/tfreplknet
PaddlePaddle	https://github.com/BR-IDL/PaddleViT/tree/develop/image_classification/RepLKNet
...

More re-implementations are welcomed.

Use our efficient large-kernel convolution with PyTorch

We have released an example for PyTorch. Please check setup.py and depthwise_conv2d_implicit_gemm.py (a replacement of torch.nn.Conv2d) in https://github.com/MegEngine/cutlass/tree/master/examples/19_large_depthwise_conv2d_torch_extension.

1. unzip cutlass.zip, enter the directory. This version of cutlass works fine with our large-kernel implementation and multiple python versions. You may alternatively use the cutlass branch maintained by the MegEngine team (clone https://github.com/MegEngine/cutlass), but you may need to be more careful with your python version (see this issue).
2. cd examples/19_large_depthwise_conv2d_torch_extension
3. ./setup.py install --user. If you get errors, check your CUDA_HOME.
4. A quick check: python depthwise_conv2d_implicit_gemm.py
5. Add WHERE_YOU_CLONED_CUTLASS/examples/19_large_depthwise_conv2d_torch_extension into your PYTHONPATH so that you can from depthwise_conv2d_implicit_gemm import DepthWiseConv2dImplicitGEMM anywhere. Then you may use DepthWiseConv2dImplicitGEMM as a replacement of nn.Conv2d.
6. export LARGE_KERNEL_CONV_IMPL=WHERE_YOU_CLONED_CUTLASS/examples/19_large_depthwise_conv2d_torch_extension so that RepLKNet will use the efficient implementation. Or you may simply modify the related code (get_conv2d) in replknet.py.

It should work with a wide range of GPUs and PyTorch/CUDA versions. We suggest you try first and check the environments only if you get any errors. Our latest testes used both

1. Ubuntu 18.04 + CUDA 11.3 + nvcc 11.3 + cudnn 8.2.0 + python 3.8.12 + pytorch 1.10 + gcc 7.3.0 + nccl 2.10.3 + NVIDIA driver 450.102.04 + V100 and A100 GPUs
2. Ubuntu 18.04 + CUDA 10.2 + nvcc 10.0 + cudnn 7.6.5 + python 3.6.9 + pytorch 1.9 + gcc 7.5.0 + nccl 2.7.8 + NVIDIA driver 460.32.03 + 2080Ti and V100 GPUs

It is reported (see here) that a python version mismatch may result in an error (forward_fp32.cu(212): error: more than one instance of constructor "cutlass::Tensor4DCoord::Tensor4DCoord" ... or cutlass/include/cutlass/fast_math.h(741): error: no suitable conversion function from "__half" to "float" exists). Please upgrade or downgrade your python. We sincerely thank @sleeplessai and @ewrfcas for sharing their experience.

Our implementation mentioned in the paper has been integrated into MegEngine. The engine will automatically use it. If you would like to use it in other frameworks like Tensorflow, you may need to compile our released cuda sources (the *.cu files in the above example should work with other frameworks) and use some tools to load them, just like cutlass and torch.utils.cpp_extension in the PyTorch example. Would be appreciated if you could share with us your experience.

You may refer to the MegEngine source code: https://github.com/MegEngine/MegEngine/tree/8a2e92bd6c5ac02807b27d174dce090ee391000b/dnn/src/cuda/conv_bias/chanwise.

Pull requests (e.g., better or other implementations or implementations on other frameworks) are welcomed.

Catalog

• [x] Model code
• [x] PyTorch pretrained models
• [x] PyTorch large-kernel conv impl
• [x] PyTorch training code
• [x] PyTorch downstream models
• [x] PyTorch downstream code
• [x] A script to visualize the ERF
• [x] How to obtain the shape bias

Results and Pre-trained Models

ImageNet-1K Models

name	resolution	ImageNet-1K acc	#params	FLOPs	ImageNet-1K pretrained model
RepLKNet-31B	224x224	83.5	79M	15.3G	Google Drive, Baidu
RepLKNet-31B	384x384	84.8	79M	45.1G	Google Drive, Baidu

ImageNet-22K Models

name	resolution	ImageNet-1K acc	#params	FLOPs	22K pretrained model	1K finetuned model
RepLKNet-31B	224x224	85.2	79M	15.3G	Google Drive, Baidu	Google Drive, Baidu
RepLKNet-31B	384x384	86.0	79M	45.1G	-	Google Drive, Baidu
RepLKNet-31L	384x384	86.6	172M	96.0G	Google Drive, Baidu	Google Drive, Baidu

MegData-73M Models

name	resolution	ImageNet-1K acc	#params	FLOPs	MegData-73M pretrained model	1K finetuned model
RepLKNet-XL	320x320	87.8	335M	128.7G	Google Drive, Baidu	Google Drive, Baidu

Evaluation

For RepLKNet-31B/L with 224x224 or 384x384, we use the "IMAGENET_DEFAULT_MEAN/STD" for preprocessing (see here). For examples,

python -m torch.distributed.launch --nproc_per_node=8 main.py --model RepLKNet-31B --batch_size 32 --eval True --resume RepLKNet-31B_ImageNet-1K_224.pth --input_size 224

or

python -m torch.distributed.launch --nproc_per_node=8 main.py --model RepLKNet-31L --batch_size 32 --eval True --resume RepLKNet-31L_ImageNet-22K-to-1K_384.pth --input_size 384

For RepLKNet-XL, please note that we used mean=[0.5,0.5,0.5] and std=[0.5,0.5,0.5] for preprocessing on MegData73M dataset as well as finetuning on ImageNet-1K. This mean/std setting is also referred to as "IMAGENET_INCEPTION_MEAN/STD" in timm, see here. Add --imagenet_default_mean_and_std false to use this mean/std setting (see here). As noted in the paper, we did not use small kernels for re-parameterization.

python -m torch.distributed.launch --nproc_per_node=8 main.py --model RepLKNet-XL --batch_size 32 --eval true --resume RepLKNet-XL_MegData73M_ImageNet1K.pth --imagenet_default_mean_and_std false --input_size 320

To verify the equivalency of Structural Re-parameterization (i.e., the outputs before and after structural_reparam), add --with_small_kernel_merged true.

Training

You may use multi-node training on a SLURM cluster with submitit. Please install:

pip install submitit

If you have limited GPU memory (e.g., 2080Ti), use --use_checkpoint true to save GPU memory.

Pretrain RepLKNet-31B on ImageNet-1K

Single machine (note --update_freq 4):

python -m torch.distributed.launch --nproc_per_node=8 main.py --model RepLKNet-31B --drop_path 0.5 --batch_size 64 --lr 4e-3 --update_freq 4 --model_ema true --model_ema_eval true --data_path /path/to/imagenet-1k --warmup_epochs 10 --epochs 300 --output_dir your_training_dir

Four machines (note --update_freq 1):

python run_with_submitit.py --nodes 4 --ngpus 8 --model RepLKNet-31B --drop_path 0.5 --batch_size 64 --lr 4e-3 --update_freq 1 --model_ema true --model_ema_eval true --data_path /path/to/imagenet-1k --warmup_epochs 10 --epochs 300 --job_dir your_training_dir

In the following, we only present multi-machine commands. You may train with a single machine in a similar way.

Finetune the ImageNet-1K-pretrained (224x224) RepLKNet-31B with 384x384

python run_with_submitit.py --nodes 4 --ngpus 8 --model RepLKNet-31B --drop_path 0.8 --input_size 384 --batch_size 32 --lr 4e-4 --epochs 30 --weight_decay 1e-8 --update_freq 1 --cutmix 0 --mixup 0 --finetune RepLKNet-31B_ImageNet-1K_224.pth --model_ema true --model_ema_eval true --data_path /path/to/imagenet-1k --warmup_epochs 1 --job_dir your_training_dir --layer_decay 0.7

Pretrain RepLKNet-31B on ImageNet-22K

python run_with_submitit.py --nodes 16 --ngpus 8 --model RepLKNet-31B --drop_path 0.1 --batch_size 32 --lr 4e-3 --update_freq 1 --warmup_epochs 5 --epochs 90 --data_set image_folder --nb_classes 21841 --disable_eval true --data_path /path/to/imagenet-22k --job_dir /path/to/save_results

Finetune 22K-pretrained RepLKNet-31B on ImageNet-1K (224x224)

python run_with_submitit.py --nodes 2 --ngpus 8 --model RepLKNet-31B --drop_path 0.2 --input_size 224 --batch_size 32 --lr 4e-4 --epochs 30 --weight_decay 1e-8 --update_freq 1 --cutmix 0 --mixup 0 --finetune RepLKNet-31B_ImageNet-22K.pth --model_ema true --model_ema_eval true --data_path /path/to/imagenet-1k --warmup_epochs 1 --job_dir your_training_dir --layer_decay 0.7

Finetune 22K-pretrained RepLKNet-31B on ImageNet-1K (384x384)

python run_with_submitit.py --nodes 4 --ngpus 8 --model RepLKNet-31B --drop_path 0.3 --input_size 384 --batch_size 16 --lr 4e-4 --epochs 30 --weight_decay 1e-8 --update_freq 1 --cutmix 0 --mixup 0 --finetune RepLKNet-31B_ImageNet-22K.pth --model_ema true --model_ema_eval true --data_path /path/to/imagenet-1k --warmup_epochs 1 --job_dir your_training_dir --layer_decay 0.7 --min_lr 3e-4

Pretrain RepLKNet-31L on ImageNet-22K

python run_with_submitit.py --nodes 16 --ngpus 8 --model RepLKNet-31L --drop_path 0.1 --batch_size 32 --lr 4e-3 --update_freq 1 --warmup_epochs 5 --epochs 90 --data_set image_folder --nb_classes 21841 --disable_eval true --data_path /path/to/imagenet-22k --job_dir /path/to/save_results

Finetune 22K-pretrained RepLKNet-31L on ImageNet-1K (384x384)

python run_with_submitit.py --nodes 4 --ngpus 8 --model RepLKNet-31L --drop_path 0.3 --input_size 384 --batch_size 16 --lr 4e-4 --epochs 30 --weight_decay 1e-8 --update_freq 1 --cutmix 0 --mixup 0 --finetune RepLKNet-31L_ImageNet-22K.pth --model_ema true --model_ema_eval true --data_path /path/to/imagenet-1k --warmup_epochs 1 --job_dir your_training_dir --layer_decay 0.7 --min_lr 3e-4

Semantic Segmentation and Object Detection

We use MMSegmentation and MMDetection frameworks. Just clone MMSegmentation or MMDetection, and

1. Put segmentation/replknet.py into mmsegmentation/mmseg/models/backbones/ or mmdetection/mmdet/models/backbones/. The only difference between segmentation/replknet.py and replknet.py is the @BACKBONES.register_module.
2. Add RepLKNet into mmsegmentation/mmseg/models/backbones/__init__.py or mmdetection/mmdet/models/backbones/__init__.py. That is ... from .replknet import RepLKNet __all__ = ['ResNet', ..., 'RepLKNet']
3. Put segmentation/configs/*.py into mmsegmentation/configs/replknet/ or detection/configs/*.py into mmdetection/configs/replknet/
4. Download and use our weights. For examples, to evaluate RepLKNet-31B + UperNet on Cityscapes python -m torch.distributed.launch --nproc_per_node=8 tools/test.py configs/replknet/RepLKNet-31B_1Kpretrain_upernet_80k_cityscapes_769.py RepLKNet-31B_ImageNet-1K_UperNet_Cityscapes.pth --launcher pytorch --eval mIoU or RepLKNet-31B + Cascade Mask R-CNN on COCO python -m torch.distributed.launch --nproc_per_node=8 tools/test.py configs/replknet/RepLKNet-31B_22Kpretrain_cascade_mask_rcnn_3x_coco.py RepLKNet-31B_ImageNet-22K_CascMaskRCNN_COCO.pth --eval bbox --launcher pytorch
5. Or you may finetune our released pretrained weights (see the tips below about the batch size and