Vision Toolbox

Backbones, necks, and useful modules for Vision tasks.

Special thanks to MLDA@EEE Lab for providing GPU resources to train the models.

Installation

Install PyTorch and torchvision from conda. Then install this GitHub repo directly

conda install pytorch torchvision cudatoolkit=11.3 -c pytorch
pip install git+https://github.com/gau-nernst/vision-toolbox.git

To update

pip install -U --force-reinstall --no-deps git+https://github.com/gau-nernst/vision-toolbox.git

Usage

from vision_toolbox import backbones

model = backbones.cspdarknet53(pretrained=True)
model(inputs)                       # last feature map, stride 32
model.get_feature_maps(inputs)      # list of 5 feature maps, stride 2, 4, 8, 16, 32
model.out_channels_list             # tuple of output channels, corresponding to each feature map

For object detection, usually only the last 4 feature maps are used. It is the responsibility of the user to select the last 4 feature maps from .get_feature_maps()

outputs = model.get_feature_maps(inputs)[-4:]   # last 4 feature maps

Backbones with ported weights

• ViT: AugReg (Google), DeiT/DeiT3 (Facebook), SigLIP (Google)
• MLP-Mixer
• CaiT
• Swin
• ConvNeXt and ConvNeXt-V2

Backbones trained by me

Implemented backbones:

• Darknet
• VoVNet

ImageNet pre-training

Download ImageNet Challenge here: https://www.kaggle.com/c/imagenet-object-localization-challenge/

kaggle competitions download -c imagenet-object-localization-challenge
unzip imagenet-object-localization-challenge.zip -d ImageNet
tar -xf ImageNet/imagenet_object_localization_patched2019.tar.gz -C ImageNet
python ./scripts/imagenet.py --val_solution_path ./ImageNet/LOC_val_solution.csv --val_image_dir ./ImageNet/ILSVRC/Data/CLS-LOC/val

To create WebDataset shards

pip install webdataset        # install webdataset
python ./scripts/wds.py --data_dir ./ImageNet/ILSVRC/Data/CLS-LOC/train --save_dir ./ImageNet/webdataset/train --shuffle True --name imagenet-train
python ./scripts/wds.py --data_dir ./ImageNet/ILSVRC/Data/CLS-LOC/val --save_dir ./ImageNet/webdataset/val --shuffle False --name imagenet-val

There should be 147 shards of training set, and 7 shards of validation set. Each shard is 1GB.

Reference recipes:

• https://github.com/pytorch/vision/blob/main/references/classification/train.py
• https://pytorch.org/blog/how-to-train-state-of-the-art-models-using-torchvision-latest-primitives/
• Ross Wightman, ResNet strikesback: https://arxiv.org/abs/2110.00476

Recipe used in this repo:

• Optimizer: SGD
• Epochs: 100
• Learning rate: 0.5 for batch size 1024, linear scaling for other batch sizes
• Weight decay: 2e-5
• Batch size: 1024, unless stated otherwise (for larger models)
• LR schedule: Linear warmup for 5 epochs, then cosine annealing
• Augmentations: (1) Random Resized Crop, (2) Trivial Augmentation, (3) Random Erasing (p=0.1), (4) CutMix (alpha=1.0), and (5) MixUp (alpha=0.2). For each batch, either CutMix or MixUp is applied, but not both.
• Label smoothing = 0.1
• Val resize: 232, Val crop: 224
• Train resized crop: 176 (FixRes)
• Mixed-precision training
• For small models, Random Erasing, CutMix, Mixup, and Label smoothing are removed (e.g. Darknet-19, VoVNet-19)
• For large models, FixRes is removed -> train resized crop is 224 (e.g. VoVNet-99)

Note: All hyperparameters are adopted from torchvision's recipe, except number of epochs (600 in torchvision's vs 100 in mine). Since the training time is shorter, augmentations should be reduced. Model EMA is not used.

PyTorch Lightning is used to train the models (see classifier.py). The easiest way to run training is to use Lightning CLI with a config file (see below, jsonargparse[signatures] is required). Note that PyTorch Lightning is not required to create, run, and load the models. Some config files are provided in this repo.

git clone https://github.com/gau-nernst/vision-toolbox.git
cd vision-toolbox
pip install pytorch-lightning jsonargparse[signatures]    # dependencies for Lightning CLI
python train.py fit --config config.yaml
python train.py fit --config config.yaml --model.backbone cspdarknet53    # change backbone to train
python train.py fit --config config.yaml --config config_wds.yaml         # use webdataset

Darknet

Paper: [YOLOv2] [YOLOv3] [CSPNet] [YOLOv4]

• Darknet-{19,53}
• CSPDarknet-53
• Darknet-YOLOv5{n,s,m,l,x}

*FLOPS is measured with (1,3,224,224) input.

For reference, official Darknet models (Sources: [pjreddie's website] [WongKinYiu's CSP GitHub repo]):

• Darknet-19: top-1 72.9%, 7.29G FLOPS
• Darknet-53: top-1 77.2%, 41.57M params, 18.57 FLOPS
• CSPDarknet-53: top-1 77.2%, 27.61M params, 13.07 FLOPS
• These models use 256x256 image, thus their FLOPS are slightly higher
• The 1000-class classification head is probaby included in their Parameters and FLOPS count, resulting in slightly higher numbers.

Convert YOLOv5 backbone weights to use with Ultralytics' repo (WIP)

python scripts/convert_yolov5_weights.py {weights_from_this_repo.pth} {save_path.pth}

The weights will be renamed to be compatiable with Ultralytics' repo. Note that the converted .pth file only contains the renamed state dict, up to model.8 (the backbone part, without the SPPF layer). You will need to modify their train script to treat the loaded file as a state dict, instead of a dictionary with key model containing the model object.

I haven't tested training a full YOLOv5 object detector with the converted weights, so this function is not guaranteed to work correctly.

VoVNet

Paper: [VoVNetV1] [VoVNetV2]

• VoVNet-(27-slim,39,57)
• VoVNet-{19-slim,19,39,57,99}-ese

*FLOPS is measured with (1,3,224,224) input.

torchvision

Some torchvision classification models are ported to use with the toolbox. They can output multiple feature map levels. torchvision>=0.11.0 is required since its feature extraction API is used to do this.

• For MobileNet and EfficientNet models, intermediate outputs are taken after the first 1x1 conv expansion layer of the strided MBConv block. See Section 6.2 of MobileNetv2 paper and Section 6.3 of MobileNetv3 paper.
• To use weights from the new PyTorch training recipe (which are significantly better), go to torchvision's prototype directory and copy weights URLs (labelled as ImageNet1K_V2) from their respective models' files.

ResNet:

• ResNet-{18,34,50,101,152}
• ResNeXt-{50,101}
• Wide ResNet-{50_2,101_2}

RegNet:

• RegNetX-{400MF,800MF,1.6GF,3.2GF,8GF,16GF,32GF}
• RegNetY-{400MF,800MF,1.6GF,3.2GF,8GF,16GF,32GF} (with SE)

MobileNet:

• MobileNetV2
• MobileNetV3-{large,small}

EfficientNet:

• EfficientNet-{B0-B7}

Necks

• FPN
• SemanticFPN
• PAN
• BiFPN

Implementation notes

• FPN: no batch norm and activation are used in lateral connections. Output convolutions (with batch norm and ReLU) are inside the top-down path.

  P5 = lateral_conv5(C5)
  P4 = out_conv4(lateral_conv4(C4) + upsample(P5))
  P3 = out_conv4(lateral_conv4(C3) + upsample(P4))

• BiFPN: weighted fusion is not implemented. Normal 3x3 convolutions (with batch norm and ReLU) are used by default.