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.
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% |
- 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
- ptflops
- opencv = 4.50
- torchvision = 0.9.0
- tensorboard >= 1.15 (optional)
- apex
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}
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 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 '''
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.