tianxintony/gsnet

Lightweight Object Detection Network for UAV Image by a Zero-Shot NAS Method.

GSNet

Introduction

Our based code is forked from ZenNAS. We modify the code to make it suitable for UAV embedding embedded platforms.

We mainly made the following changes:

• redesign a new lightweight search space, GhostShuffle unit(GS unit), which further reduce the params and FLOPs of network.
• add constraints on params, FLOPs ,layers and MAC, and searches for the optimal network GSmodel-L based on the ZenScore scores.

Experimental results

We tested the modified code and verified its correctness. The results are as follows:

We used apex with mixed precision to complete the training within 5 days on 4 2080Ti GPUs, and the results are consistent with the paper.

dataset	paper model params	paper model FLOPs	mAP
VisDrone	6.61M	11.20M	14.92%
UAV-OUC-DET	6.61M	11.20M	8.38%

Reproduce Paper Experiments

System Requirements

• PyTorch = 1.8.0, Python = 3.7.9, CUDA=10.2
• By default, ImageNet dataset is stored under ~/data/imagenet; CIFAR-10/CIFAR-100 is stored under ~/data/pytorch_cifar10 or ~/data/pytorch_cifar100
• Pre-trained parameters are cached under ~/.cache/pytorch/checkpoints/zennet_pretrained

Package Requirements

• ptflops
• opencv = 4.50
• torchvision = 0.9.0
• tensorboard >= 1.15 (optional)
• apex

Evaluate pre-trained models on ImageNet and CIFAR-10/100

To evaluate the pre-trained model on ImageNet using GPU 0:

cd GSNet/src/zen_nas/

python test.py --fp16 --gpu 0 --arch ${zennet_model_name}

where ${zennet_model_name} should be replaced by a valid ZenNet model name. The complete list of model names can be found in the 'Pre-trained Models' section.

To evaluate the pre-trained model on CIFAR-10 or CIFAR-100 using GPU 0:

python test_cifar.py --dataset cifar10 --gpu 0 --arch ${zennet_model_name}

To create a ZenNet in your python code:

gpu=0 model = ZenNet.get_ZenNet(opt.arch, pretrained=True) torch.cuda.set_device(gpu) torch.backends.cudnn.benchmark = True model = model.cuda(gpu) model = model.half() model.eval()

Searching on CIFAR-10/100

Searching for CIFAR-10/100 models with budget params < 1M , using different zero-shot proxies:

'''bash scripts/Flops_NAS_cifar_params1M.sh scripts/GradNorm_NAS_cifar_params1M.sh scripts/NASWOT_NAS_cifar_params1M.sh scripts/Params_NAS_cifar_params1M.sh scripts/Random_NAS_cifar_params1M.sh scripts/Syncflow_NAS_cifar_params1M.sh scripts/TE_NAS_cifar_params1M.sh scripts/Zen_NAS_cifar_params1M.sh '''

Designing the object detection network, GSNet.

After the search stage, we got the structure of backbone. According to the result, we construct a backbone network and a Ghost-PAN feature fusion module and detection heads are added to complete the design of the lightweight object detection network, GSNet. Of course, there is a easier method to work on object detection that is, we can modify the backbone based on nanodet. In addition, we have provided a reference and added a file GSnetwork, the ghostshuffle.py is the whole network structure of GSNet according our search result.