StitchingNet Fine-Tuning 및 모델 경량화

프로젝트 개요

본 프로젝트는 공개된 StitchingNet 데이터셋을 사용하여 봉제 불량을 탐지하는 경량 CNN 모델들을 Fine-Tuning 하고, 이를 라즈베리파이와 같은 엣지 디바이스 환경에서 실시간 추론할 수 있도록 모델을 압축(양자화, 채널 조정, Pruning)하는 것을 목표로 합니다.

MobileNet 시리즈를 포함한 다양한 CNN 구조를 평가하고, 정확도를 유지하면서도 모델 크기와 연산량을 최소화하기 위해 PTQ/QAT 양자화, 알파(Width Multiplier) 조절 등의 기법을 적용하였습니다. 또한 Pruning 및 Knowledge Distillation 기법들을 추가적으로 적용해볼 예정입니다.

데이터셋 정보

Kaggle에 공개된 StitchingNet 데이터셋을 사용합니다. 이 데이터셋은 11가지 클래스(정상과 10가지 불량 유형)로 구분된 총 14,565장의 봉제 이미지(224×224 크기)를 포함하고 있으며, 다양한 원단과 실의 색상을 고려하여 만들어졌습니다.

학습을 위해 데이터셋을 훈련, 검증, 테스트 세트로 나누고 다양한 데이터 증강 기법을 적용하였습니다.

모델 경량화 기법

다음의 경량화 기법을 적용하여 엣지 디바이스에서 실시간 추론이 가능한 모델을 제작하였습니다.

• 양자화(Quantization):
  • PTQ(Post-Training Quantization): 추가 학습 없이 FP32 모델을 INT8로 변환하여 속도 향상.
  • QAT(Quantization-Aware Training): Fake Quantization을 적용하여 학습함으로써 양자화로 인한 정확도 손실 최소화.
• Width Multiplier 조정: 각 레이어 채널 수를 축소해 FLOPs 및 모델 크기를 조정.

추가적으로 적용을 고려하고 있는 기법은 다음과 같습니다다. - Pruning: 중요하지 않은 채널을 제거하는 Pruning 기법을 활용하여 모델 크기 및 연산량 감소. - Knowledge Distillation (KD): 성능이 좋은 큰 모델(Teacher)의 지식을 작은 모델(Student)로 이전하여 성능 향상.

데이터 증강 기법

모델의 일반화를 높이기 위해 아래의 데이터 증강 버전을 설계하고 테스트하였습니다.

• Version 0: 원본 이미지 사용
• Version 1: 기하학적 변형 (상하좌우 반전, 회전)
• Version 2: 기하학적 변형 + 색상 변화(Hue/Saturation 변경, 컬러 지터)
• Version 3: 기하학적 변형 + 색상 변화 + 블러 및 노이즈(가우시안 블러 및 노이즈)

실험 결과 Version 1과 Version 2가 가장 우수한 성능을 나타냈습니다.

프로젝트 사용 방법

설치

참고: 본 프로젝트는 Python 3.10.16 환경에서 개발 및 테스트되었습니다.

GitHub 저장소를 클론한 후 다음 명령어로 필요한 라이브러리를 설치합니다.

```bash git clone https://github.com/KimJunWon98/StitchingNet-Finetuning.git cd StitchingNet-Finetuning

pip install -r requirements.txt ```

StitchingNet 데이터셋을 다운로드한 후, ./dataset/StitchingNet 폴더에 압축을 해제합니다.

dataset/ └── StitchingNet ├── A. Cotton-Poly ├── B. Linen-Poly ├── C. Denim-Poly ├── D. Velveteen-Poly ├── E. Polyester-Poly ├── F. Satin-Core ├── G. Chiffon-Poly ├── H. Nylon-Core ├── I. Jacquard-Poly ├── J. Oxford-Core └── K. Polyester (coated)-Core

학습 실행 예시

모델을 학습하려면 다음 명령어를 사용합니다.

bash python main.py --model mobilenetv1_100 --augment 1 --epochs 100 --batch-size 32 --patience 5

양자화 적용 예시

양자화를 적용하여 학습하려면 다음 명령어를 사용합니다.

• 양자화: bash python quantization-train.py --model mobilenet_v2 --augment 2 --epochs 100 --batch-size 32 --patience 5

Fine-tuning 결과

모델의 크기와 성능에 대한 결과는 아래 표를 참고하세요. (증강 및 양자화를 적용하지 않은 결과입니다.)

| Model Name | total FLOPs | FLOPs (% of MobileNetV1) | Test Accuracy (%) | Test Loss | model size (MB) | total params | |:----------------------------:|:---------------:|:----------------------------:|:---------------------:|:-------------:|:-------------------:|:----------------:| | mobilenetv3small050 | 27,760,784 | 4.68% | 97.162% | 0.08815 | 2.2106 | 579,499 | | lcnet_050 | 45,998,976 | 7.76% | 98.764% | 0.05287 | 2.3420 | 613,947 | | repghostnet_050 | 47,806,496 | 8.06% | 98.261% | 0.06630 | 3.9942 | 1,047,059 | | mobilenetv3small075 | 48,897,336 | 8.25% | 97.574% | 0.08905 | 3.9221 | 1,028,147 | | mobilenetv3small100 | 62,589,728 | 10.55% | 99.497% | 0.02862 | 5.8332 | 1,529,131 | | mobilenetv4convsmall_050 | 65,866,112 | 11.11% | 99.268% | 0.03905 | 3.7080 | 972,043 | | tinynet_e | 92,972,848 | 15.68% | 98.993% | 0.04775 | 2.9604 | 776,063 | | lcnet_075 | 97,522,112 | 16.45% | 98.902% | 0.05426 | 4.1633 | 1,091,379 | | repghostnet_080 | 103,419,948 | 17.44% | 98.856% | 0.06593 | 7.6754 | 2,012,059 | | tinynet_d | 108,160,480 | 18.24% | 99.314% | 0.02845 | 4.0876 | 1,071,537 | | mobilenetv2_050 | 114,129,200 | 19.25% | 95.744% | 0.14250 | 2.6770 | 701,771 | | tinynet_c | 146,790,496 | 24.75% | 99.268% | 0.03166 | 4.5407 | 1,190,325 | | ghostnet_100 | 147,008,016 | 24.79% | 99.222% | 0.04460 | 14.9368 | 3,915,599 | | repghostnet_100 | 152,463,280 | 25.71% | 99.497% | 0.03870 | 10.7067 | 2,806,695 | | lcnet_100 | 159,502,336 | 26.90% | 98.032% | 0.08151 | 6.4350 | 1,686,891 | | ghostnetv2_100 | 176,307,952 | 29.73% | 98.764% | 0.04526 | 18.6539 | 4,889,999 | | mobilenetv4convsmall | 188,715,520 | 31.82% | 99.634% | 0.02742 | 9.5639 | 2,507,115 | | mobilenetv3_rw | 237,368,096 | 40.03% | 99.085% | 0.05632 | 16.0714 | 4,213,009 | | mobilenetv3large100 | 237,369,920 | 40.03% | 99.588% | 0.02299 | 16.0832 | 4,216,123 | | repghostnet_130 | 244,782,348 | 41.28% | 99.314% | 0.03422 | 16.0664 | 4,211,703 | | mixnet_s | 250,217,856 | 42.20% | 99.542% | 0.02925 | 9.9736 | 2,614,513 | | ghostnetv2_130 | 281,690,296 | 47.50% | 99.268% | 0.05071 | 29.3484 | 7,693,495 | | tinynet_b | 287,347,088 | 48.46% | 99.497% | 0.02773 | 9.3981 | 2,463,653 | | mnasnet_100 | 324,125,952 | 54.66% | 98.902% | 0.04141 | 11.8881 | 3,116,403 | | mobilenetv2_100 | 332,961,632 | 56.15% | 98.947% | 0.05497 | 8.5372 | 2,237,963 | | spnasnet_100 | 344,324,928 | 58.07% | 99.268% | 0.03387 | 12.0343 | 3,154,707 | | mixnet_m | 354,094,048 | 59.71% | 99.588% | 0.02790 | 13.3297 | 3,494,289 | | ghostnetv2_160 | 414,088,472 | 69.83% | 99.725% | 0.03040 | 42.4415 | 11,125,789 | | tinynet_a | 471,260,448 | 79.47% | 99.725% | 0.01753 | 18.7724 | 4,921,063 | | mobilenetv2_110d | 472,868,000 | 79.74% | 99.451% | 0.02706 | 12.3963 | 3,249,611 | | repghostnet_200 | 537,274,080 | 90.60% | 99.405% | 0.04581 | 32.5364 | 8,529,215 | | mixnet_l | 575,452,624 | 97.04% | 99.497% | 0.02967 | 22.1602 | 5,809,159 | | mobilenetv1_100 | 592,991,232 | 100.00% | 99.634% | 0.01746 | 12.2767 | 3,218,251 | | mobilenetv2_140 | 628,378,464 | 105.97% | 99.085% | 0.03273 | 16.5386 | 4,335,499 | | mobilenetv2_120d | 727,481,408 | 122.68% | 99.268% | 0.03854 | 17.4112 | 4,564,235 | | mobilenetv3large150d | 821,242,720 | 138.49% | 99.497% | 0.02623 | 50.9468 | 13,355,411 | | mobilenetv4convmedium | 838,036,864 | 141.32% | 99.405% | 0.02861 | 32.2289 | 8,448,603 | | mixnet_xl | 931,662,848 | 157.11% | 99.634% | 0.01866 | 39.5839 | 10,376,675 | | mobilenetv4hybridmedium| 952,689,024 | 160.66% | 99.680% | 0.01663 | 37.4136 | 9,807,739 | | resnet18 | 1,818,559,488 | 306.68% | 99.451% | 0.02435 | 42.6565 | 11,182,155 |

Quantization 결과

모델의 크기, 성능, 추론 속도는 아래 표를 참고하세요. (데이터 증강을 적용하지 않은 결과입니다.)
※ 추론 속도는 CPU 사양에 따라 달라질 수 있으며, 본 실험은 AMD Ryzen 5 9600X에서 측정되었습니다.

| Model | Quantization | Test Accuracy | Test Loss | Inference Time (ms) | |:---------------------------------------:|:----------------:|:-----------------:|:-------------:|:-----------------------:| | Mobilenet V2 (α = 1) | X | 99.04% | 0.0400 | 7.779 | | Mobilenet V2 (α = 1) | PTQ | 96.52% | 0.1202 | 2.779 | | Mobilenet V2 (α = 1) | QAT | 98.17% | 0.0856 | 3.150 | | Mobilenet V3 large (α = 1) | X | 99.36% | 0.0506 | 8.560 | | Mobilenet V3 large (α = 1) | PTQ | 71.85% | 1.6623 | 3.543 | | Mobilenet V3 large (α = 1) | QAT | 93.64% | 0.2410 | 3.415 |