YOLOv8 모델 구조 수정 방법

처음 YOLOv8 모델을 개선하려고 했을 때,
자료가 너무 부족해서 어느 파일의 어떤 부분을 고쳐야 하는지 찾기 힘들었다 😂

그래서, 어떻게 수정해야 하는지를 공유하고자 한다.

이번 포스트에서는 새로운 모듈을 추가하고, 모델의 구조를 수정하는 방법에 대해 설명할 것이다.

---

1. YOLOv8 패키지 클론

먼저, YOLOv8을 개발한 ultralytics의 깃허브에서 코드를 클론해준다.

https://github.com/ultralytics/ultralytics

코드를 보면 ultralytics 안에 또다른 ultratlytics 폴더가 있을 것이다.
들어가보자.

여기서 우리가 주목해야 할 곳은 바로 cfg 폴더와 nn 폴더이다.

---

2. 사용할 모듈 정의

nn/modules 폴더에 들어가보자.
파일들을 열어보면, 여러 모듈들이 정의되어있는 것을 볼 수 있다.

우리도 여기에 모듈을 추가해줄 것이다.

1) 먼저, 새로운 파일을 만들어준다.

나는 coordatt.py라는 파일을 만들어주었다.

2) 그 다음, 만들어준 파일에 pytorch 프레임워크로 새로운 모듈 클래스를 정의해준다.

나는, 논문 저자가 pytorch로 구현해놓은 coordinate attention 클래스를 사용했다.
CoordAtt라는 이름으로 클래스를 정의하였다.

사용한 코드: https://github.com/houqb/CoordAttention/blob/main/coordatt.py

* 여기서 주의해야 할 점은 이 파일 내에
__all__ 변수에 ('모듈 이름')을 써주어야 한다는 것이다!

이런식으로 하면 된다.

3) 이제, 이 모듈이 추가되었다는 것을 __init__.py에 알려주어야 한다.

먼저, nn/modules 폴더 내에 있는 __init__.py에 들어간다.
(modules 폴더 바깥에도 __init__.py가 있는데 헷갈리면 안된다.)

그 다음, 1)에서 추가해준 파일에서 클래스를 불러온다.

from .coordatt import CoordAtt

그 다음, __all__ 튜플에 클래스 이름을추가해준다.

__all__ = ('Conv', 'Conv2', 'LightConv', 'RepConv', 'DWConv', 'DWConvTranspose2d', 
...
(중략)
...
, 'CoordAtt')

이렇게 하면, 모듈 정의는 끝난다!

---

3. 모델 구조에 모듈을 추가

이제 정의한 모듈을 모델 구조에 추가해주는 일만 남았다.

1) nn/modules 폴더를 나와서 ultratlycis/ultralytics/cfg/models/v8 폴더에 들어간다.

2) yolov8.yaml 파일을 복사, 붙여넣기 해준다.

yolov8.yaml이 가장 기초가 되는 모델 구조 파일이다.

잠깐 이 파일에 대해 설명을 하자면,
각 줄이 하나의 레이어인 것이다.

각 줄의
첫 번째 요소는 이 레이어가 어디에 이어졌는지를 명시한다.
-1이라면 바로 앞 레이어이와 이어진 것이고
[-1, 12]라면, 바로 앞 레이어와 12번째 레이어를 연결하는 것이다.
두 번째 요소는 이 레이어를 몇 번 반복할 것인지를 명시한다.
세 번째 요소는 어떤 모듈을 사용할 것인지를 명시한다.
네 번째 요소는 이 레이어를 사용할 때 들어갈 인자값을 명시한다.

3) 원하는 부분에 모듈을 추가해준다.

아까 2.2)에서 정의한 모듈을 사용한다면, 그 클래스 이름을 사용할 것이다.

노란색으로 밑줄 친 부분이 내가 새로 추가한 부분이다.

16번째 레이어로 아까 정의한 CoordAtt 모듈을 추가했다.

이렇게 모듈을 추가할 때 신경써야 할 부분은 두 가지가 있다.

1) 추가한 모듈에 들어갈 인자값을 신경써야 한다.

2) concat 레이어와 detect 레이어는 특정 레이어층을 연결하는데 이 레이어 숫자에 신경써야 한다.

예를 들어, 18번째 자리에 있는 concat 레이어가 8번째 레이어와 17번째 레이어를 연결하고 있었는데,
16번째 자리에 레이어가 하나 추가되었다면, 17번째 레이어는 한 칸이 뒤로 밀릴 것이다.

그럼 이 concat 레이어는 8번째 레이어와 18번째 레이어를 연결하도록 숫자를 수정해야 할 것이다.

---

4. 수정한 모델 구조 학습시키기

이제 모든 수정 작업은 끝났다.
모델 구조를 불러와서 학습만 시키면 된다.

내가 만든 모델 구조 파일을 불러오는 것은 아주 간단하다.

model = YOLO('주소/ultralytics/ultralytics/cfg/models/v8/yolov8n-ca.yaml')

• 주의할 점은 모델 구조가 yolov8-ca.yaml 파일에 저장되어있고,
  내가 nano 모델을 사용하고 싶다면, yolov8n-ca.yaml로 불러오면 되고,
  large 모델을 사용하고 싶다면, yolov8l-ca.yaml로 불러오면 된다.

학습 시키는 것도 간단하다!

model.train(
인자값들
)

---

이렇게 해서 간단하게 yolov8 모델 구조를 내가 원하는대로 수정하는 방법을 알아보았다.

다음 포스트에서는 yolov8의 loss function을 수정하는 방법에 대해 설명할 것이다.