devlee247 / yolox-nano-model-compression Goto Github PK

인공지능 모델 압축 (경량화) 기술 기반 On-device AI를 추구하는 Nota AI와 연계하여 진행한 YOLOX-nano 모델 경량화 프로젝트입니다.

YOLOX-Nano 모델 구조 변경을 통한 YOLOX-nano 정확도(mAP) 개선

• 모델의 FLOPs와 Parameter 수는 기준 모델보다 작거나 같아야 함
• 고도화된 학습 기법을 이용한 성능 개선 금지
• Dataset : Pascal VOC 2007

모델을 구성하는 레이어, 합성곱 블록 등을 변형하여 베이스라인 모델의 parameter 개수와 GFLOPs를 초과하지 않으면서 더욱 정확도(mAP) 높은 모델을 생성

먼저, 기존의 모델에서 구조적 개선을 통한 경량화를 진행하여 최대한 성능을 유지하면서 FLOPs와 Parameters를 줄일 수 있었습니다.

여유 공간을 확보한 후, 성능 향상 기법을 적용하여 기존 모델 대비 Parameters는 약 11.1% 감소, [email protected]:.95는 약 12.4% 상승이라는 결과를 얻어냈습니다.

• 구조 개선을 통한 경량화
  • Backbone 구조 개선
  • FPN 구조 개선
• 성능 향상 기법
  • Mobile Bottleneck
  • GhostNet Convolution
  • Spatial Attention

Methods별 실험 결과

conda create -n yolox python=3.7

conda activate yolox

conda install pytorch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2 cudatoolkit=11.0 -c pytorch

apt-get install gcc

apt-get install g++

pip3 install -v -e .

pip3 install wandb

cd YOLOX/datasets
mkdir trainval
mkdir test

# train dataset
wget http://pjreddie.com/media/files/VOCtrainval_06-Nov-2007.tar
tar -xvf VOCtrainval_06-Nov-2007.tar -C ./trainval
rm -rf VOCtrainval_06-Nov-2007.tar

# val dataset
wget http://pjreddie.com/media/files/VOCtest_06-Nov-2007.tar
tar -xvf VOCtest_06-Nov-2007.tar -C ./test
rm -rf VOCtest_06-Nov-2007.tar

bash run.sh

위의 명령어로 바로 train을 진행할 수 있습니다.

shell script의 상세 구성은 다음과 같습니다.

####################### 이 부분만 수정하세요! #######################

DIR_NAME='[name0]name_of_experiments' # 템플릿 폴더 이름
BATCH_SIZE=64 # 배치 사이즈
SEED=42 # 시드 설정
EVAL_INTERVAL=10

# WANDB Setting
WANDB_LOG=true # true 값으로 설정해야 Wandb에 로그가 올라갑니다.
PROJECT_NAME='YOLOX_NANO_TRAIN_NOTA' # Wandb Project Name

################################ END ################################

DIR_NAME, BATCH_SIZE, SEED, EVAL_INTERVAL, WANDB_LOG, PROJECT_NAME을 수정하여 학습을 진행할 수 있습니다.

이 외에 블록 교체 등의 작업은 yolox/models/{DIR_NAME}에서

yolox.py, yolo_pafpn.py, yolo_head.py, darknet.py, yolox_nano_exp.py를 수정하여 진행하시면 됩니다.

이외에 Conv layer 등을 수정하려면 yolox/models/network_blocks.py를 수정하시면 됩니다.

Shell script 말고도 학습을 진행할 수 있습니다.

python tools/train.py -f exps/default/yolox_voc_nano.py -b 16 --logger wandb wandb-project <project name>

김찬혁	문태진	이인서	장상원

If you use YOLOX in your research, please cite our work by using the following BibTeX entry:

 @article{yolox2021,
  title={YOLOX: Exceeding YOLO Series in 2021},
  author={Ge, Zheng and Liu, Songtao and Wang, Feng and Li, Zeming and Sun, Jian},
  journal={arXiv preprint arXiv:2107.08430},
  year={2021}
}