참고
https://github.com/Andrewhsin/YOLO-NAS-pytorch
위 링크에서 train.py와 inference.py 파일을 가져왔다.
데이터셋 준비
COCO 혹은 YOLO 데이터 포맷으로 된 커스텀 데이터셋을 아래와 같이 준비한다.
├── Dataset
│ ├── train
│ │ ├── images
│ │ │ ├── 1.jpg
│ │ │ ├── 2.jpg
| | | ├── ....
│ │ ├── labels
│ │ │ ├── 1.txt
│ │ │ ├── 2.txt
| | | ├── ....
│ ├── val
│ │ ├── images
│ │ │ ├── 1.jpg
│ │ │ ├── 2.png
| | | ├── ....
│ │ ├── labels
│ │ │ ├── 1.txt
│ │ │ ├── 2.txt
| | | ├── ....
│ ├── test
│ │ ├── images
│ │ │ ├── 1.jpg
│ │ │ ├── 50.jpg
| | | ├── ....
│ │ ├── labels
│ │ │ ├── 1.txt
│ │ │ ├── 50.txt
| | | ├── ....커스텀 데이터셋으로 모델을 학습하기 위해 data.yaml 파일을 아래와 같이 작성한다.
names: # 클래스 이름
- class1
- class2
- class3
Dir: '/home/user/my_data' # 커스텀 데이터셋 경로
images:
test: test/images
train: train/images
val: valid/images
labels:
test: test/labels
train: train/labels
val: valid/labels학습 수행
train.py에 data.yaml 경로와 batch, epoch, 모델 등을 입력하여 학습을 수행한다.
-i, --data: path to data.yaml
-n, --name: Checkpoint dir name
-b, --batch: Training batch size
-e, --epoch: number of training epochs.
-j, --worker: Training number of workers
-m, --model: Model type (eg: yolo_nas_s, yolo_nas_m, yolo_nas_l)
-w, --weight: path to pre-trained model weight (default: coco weight)python3 train.py --data /dir/dataset/data.yaml --batch 6 --epoch 100 --model yolo_nas_m이후 아래와 같이 로그가 출력되면서 학습이 진행된다.
[2024-03-07 14:44:41] INFO - super_gradients.training.utils.sg_trainer_utils - TRAINING PARAMETERS:
- Mode: Single GPU
- Number of GPUs: 1 (1 available on the machine)
- Full dataset size: 2000 (len(train_set))
- Batch size per GPU: 16 (batch_size)
- Batch Accumulate: 1 (batch_accumulate)
- Total batch size: 16 (num_gpus * batch_size)
- Effective Batch size: 16 (num_gpus * batch_size * batch_accumulate)
- Iterations per epoch: 125 (len(train_loader))
- Gradient updates per epoch: 125 (len(train_loader) / batch_accumulate)
- Model: YoloNAS_S (19.02M parameters, 19.02M optimized)
- Learning Rates and Weight Decays:
- default: (19.02M parameters). LR: 0.0005 (19.02M parameters) WD: 0.0, (42.13K parameters), WD: 0.0001, (18.98M parameters)
[2024-03-07 14:44:41] INFO - super_gradients.training.sg_trainer.sg_trainer - Started training for 300 epochs (0/299)
[2024-03-07 14:46:34] INFO - super_gradients.common.sg_loggers.base_sg_logger - Checkpoint saved in runs/train0/RUN_20240307_144436_802501/ckpt_best.pth
[2024-03-07 14:46:34] INFO - super_gradients.training.sg_trainer.sg_trainer - Best checkpoint overriden: validation mAP@0.50: 0.09033182263374329
[2024-03-07 14:48:27] INFO - super_gradients.common.sg_loggers.base_sg_logger - Checkpoint saved in runs/train0/RUN_20240307_144436_802501/ckpt_best.pth
[2024-03-07 14:48:27] INFO - super_gradients.training.sg_trainer.sg_trainer - Best checkpoint overriden: validation mAP@0.50: 0.8522345423698425
[2024-03-07 14:50:21] INFO - super_gradients.common.sg_loggers.base_sg_logger - Checkpoint saved in runs/train0/RUN_20240307_144436_802501/ckpt_best.pth
[2024-03-07 14:50:21] INFO - super_gradients.training.sg_trainer.sg_trainer - Best checkpoint overriden: validation mAP@0.50: 0.8929525017738342SUMMARY OF EPOCH 1
├── Train
│ ├── Ppyoloeloss/loss_cls = 1.0575
│ │ ├── Epoch N-1 = 2.09 ([32m↘ -1.0325[0m)
│ │ └── Best until now = 2.09 ([32m↘ -1.0325[0m)
│ ├── Ppyoloeloss/loss_iou = 0.5165
│ │ ├── Epoch N-1 = 0.7264 ([32m↘ -0.21[0m)
│ │ └── Best until now = 0.7264 ([32m↘ -0.21[0m)
│ ├── Ppyoloeloss/loss_dfl = 0.5596
│ │ ├── Epoch N-1 = 0.723 ([32m↘ -0.1634[0m)
│ │ └── Best until now = 0.723 ([32m↘ -0.1634[0m)
│ └── Ppyoloeloss/loss = 2.1336
│ ├── Epoch N-1 = 3.5394 ([32m↘ -1.4058[0m)
│ └── Best until now = 3.5394 ([32m↘ -1.4058[0m)
└── Validation
├── Ppyoloeloss/loss_cls = 0.8775
│ ├── Epoch N-1 = 2.2145 ([32m↘ -1.337[0m)
│ └── Best until now = 2.2145 ([32m↘ -1.337[0m)
├── Ppyoloeloss/loss_iou = 0.3736
│ ├── Epoch N-1 = 0.5196 ([32m↘ -0.146[0m)
│ └── Best until now = 0.5196 ([32m↘ -0.146[0m)
├── Ppyoloeloss/loss_dfl = 0.4457
│ ├── Epoch N-1 = 0.5581 ([32m↘ -0.1124[0m)
│ └── Best until now = 0.5581 ([32m↘ -0.1124[0m)
├── Ppyoloeloss/loss = 1.6969
│ ├── Epoch N-1 = 3.2923 ([32m↘ -1.5954[0m)
│ └── Best until now = 3.2923 ([32m↘ -1.5954[0m)
├── Precision@0.50 = 0.0348
│ ├── Epoch N-1 = 0.2138 ([31m↘ -0.179[0m)
│ └── Best until now = 0.2138 ([31m↘ -0.179[0m)
├── Recall@0.50 = 0.9845
│ ├── Epoch N-1 = 0.0216 ([32m↗ 0.9629[0m)
│ └── Best until now = 0.0216 ([32m↗ 0.9629[0m)
├── Map@0.50 = 0.8522
│ ├── Epoch N-1 = 0.0903 ([32m↗ 0.7619[0m)
│ └── Best until now = 0.0903 ([32m↗ 0.7619[0m)
├── F1@0.50 = 0.067
│ ├── Epoch N-1 = 0.0393 ([32m↗ 0.0278[0m)
│ └── Best until now = 0.0393 ([32m↗ 0.0278[0m)
└── Best_score_threshold = 0.47
├── Epoch N-1 = 0.06 (↗ 0.41[0m)
└── Best until now = 0.06 (↗ 0.41[0m)
super-gradients 깃헙에는 recipe로 학습하라고 되있던데 일단 YOLO-NAS로 학습시키는건 처음이라 뭐가 뭔지도 모르겠고 커스텀 데이터셋에 대한 내용은 안나와있는거 같아서 찾아보니까 파이썬 커맨드로 학습시키는 코드가 있는걸 발견했다.
사실 YOLO 학습 자체를 별로 해본 적이 없어서 epoch나 batch 같은 수치를 얼마나 줘야하는지는 모르겠고 대충 찾아보니까 다른 코드에서는 batch 16에 epoch 50 ~ 300 정도를 쓰길래 입력해봤는데 2천장짜리 데이터셋에 한 3시간 정도 걸리는 것 같다.
아직 학습이 진행중이라 내일 결과보고 수치를 좀 더 만져봐야할 듯 하다.
'dev' 카테고리의 다른 글
| [Jetson] Jetson Orin Nano SSD에 Jetpack 설치 (SDK Manager 사용) (3) | 2024.04.18 |
|---|---|
| YOLO-NAS 로그 확인하기(텐서보드) (0) | 2024.03.08 |
| YOLOv9 예제 (0) | 2024.03.01 |
| YOLO-NAS 예제 (0) | 2024.02.29 |
| [Linux] Ubuntu 22.04 NVIDIA 드라이버 + CUDA + cuDNN 설치하기 (5) | 2024.02.29 |
