YOLOv5 이어서 학습하기 --resume 사용

YOLOv5를 학습하다가 예기지 못하게 중간에 중단 현상이 나타났다.

(사실 전적으로 본인의 잘못😭)

 

그래서 YOLOv5 학습 도중 중단되었을 때 학습을 이어나가는 방법을 정리하고자 한다. 

 

YOLOv5 이어서 학습하기

--resume을 사용하면 된다! 

 

YOLO 처럼 유명한 모델은 대부분 weight를 기록하고 저장하는 코드가 굉장히 잘 짜여 있다. 

그렇기 때문에 그런 모델을 사용할 경우에 우리는 있는 것을 잘 찾아서 사용하기만 하면 된다!

 

원래 다음과 같은 명령어를 실행해서 epoch는 100으로 설정하고  훈련시키고 있었다.

python train.py --img 640 --batch 64 --epochs 100 --data data/data.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --device 0

 

weight를 포함한 결과들이 runs/train/exp 폴더에 저장되는데

그 폴더 안 result 엑셀 파일을 통해 epoch 77까지 훈련이 완료되었고 78 도중에 끊겼다는 것을 확인하였다. 

 

다음 명령어를 사용해서 YOLO를 이어서 학습할 수 있다. 

python train.py --weights runs/train/exp/weights/last.pt --device 0 --resume

• --weights: 중간에 끊긴 모델의 weight 저장 경로로 설정하면 된다.
• --resume: 이 명령어를 사용하면 epoch를 따로 지정해주지 않고도 이어서 훈련할 수 있다.

아래가 yolov5의 train.py에 있는 내용 그대로를 가져온 것이다. 
parser.add_argument("--resume", nargs="?", const=True, default=False, help="resume most recent training")

 

위 명령어를 실행하니 다음과 같이 자동으로 이어서 훈련되는 것을 확인할 수 있었다.

Resuming training from runs/train/exp/weights/last.pt from epoch 78 to 100 total epochs

 

 

 

 

 

Reference

https://github.com/ultralytics/yolov5/blob/master/train.py

yolov5/train.py at master · ultralytics/yolov5