yolov5 custom train 5

박선후·2021년 8월 3일
0

yolov5

목록 보기
5/6

기존의 21737장의 데이터셋에 대해서 학습한 것은 만족할만한 결과를 얻지 못했다. 주로 사용할 테스트 셋은 도로주행영상인데 학습 데이터 셋에는 단순 차량 이미지만 있기 때문에, 이러한 결과가 나온 것 같다. 따라서 내가 원하는 학습 데이터 셋을 더 찾아 보았다.

AIHub

link: https://aihub.or.kr/aihub-data/autonomous/about

다음과 같이 공공데이터 사이트를 찾을 수 있었다. 여기에는 내가 원하는 학습데이터인 도로주행영상 데이터 셋을 쉽게 다운 받을 수 있었다.

xml, json 등 다양한 라벨이 포함되어 있었고, 나는 그중에 xml로 라벨링이 되어있는 데이터셋을 다운받아 전처리를 진행하였다.

직접만든 데이터 전처리 코드
https://github.com/parksh089g/data_Pre_Processing

데이터셋의 이미지 크기를 보면 HD 크기 또는 1280x720 사이즈 였다. 기존에 사용하던 yolov5x.pt(yolo pretrained weight)는 default 이미지 사이즈가 640이었기 때문에 적합하지 않다는 생각이 들었다. 최근 release를 보니 1280 사이즈에 맞는 pretrained weight가 있어 해당 가중치를 가지고 학습을 시작하였다.

177,992장 이미지 커스텀 학습 진행

  • epoch 24(추가로 진행중)/ batch-size 6 / input-size 1280 / RTX 3090 / pretrained weight yolov5x6
  • train set: 167,621
  • val set: 10,371
  • class: 'manholecover','pothole','roadcrack','none_of_the_above','Vehicle_Car','Vehicle_Bus','Vehicle_Motorcycle','Vehicle_Unknown','Pedestrian_Pedestrian','Pedestrian_Bicycle','Lane_White_Dash','Lane_White_Solid','Lane_Yellow_Dash','Lane_Yellow_Solid','Lane_Blue_Dash','Lane_Blue_Solid','TrafficLight_Red','TrafficLight_Yellow','TrafficLight_Green','TrafficLight_Arrow','TrafficLight_RedArrow','TrafficLight_YellowArrow','TrafficLight_GreenArrow','TrafficSign_Speed','TrafficSign_Else','RoadMark_StopLine','RoadMark_Crosswalk','RoadMark_Number','RoadMark_Character','RoadMarkArrow_Straight','RoadMarkArrow_Left','RoadMarkArrow_Right','RoadMarkArrow_StraightLeft','RoadMarkArrow_StraightRight','RoadMarkArrow_Uturn','RoadMarkArrow_Else','FreeSpace'

데이터셋의 양이 워낙 많고 이미지 사이즈도 크다보니 RTX 3090을 사용하고 batch-size 6 (24G 중 사용률 95%)로 설정해도 1 epoch를 학습하는데 약 10시간이 걸렸다. 약 10일 정도 학습한 결과 24에포크를 학습하였고 mAP는 기존의 학습했던 결과들과 비슷했다.(mAP: 약 0.69) 하지만 테스트 했을 때, 매우 만족스러운 결과를 보여주었다.

동영상에 대해서 detect를 진행했을 때, FPS가 30정도 나왔다.(rtx 3080 사용)

profile
공부노트

1개의 댓글

comment-user-thumbnail
2021년 11월 24일

안녕하세요, 현재 저도 AI hub에서 데이터를 받아 YOLO모델에 사용하려 하는데, 막히는 부분이 있어 질문을 드리고자 하는데 혹시 이메일을 알려주실 수 있을까요? 알려주신다면 연락드리겠습니다!

답글 달기