yolov5 환경에서 yolov4 모델 학습 및 예측 해보기

1. yolov4.yaml 생성

yolov4.yaml 경로 : yolov5/models

 

 

https://github.com/Lornatang/YOLOv4-PyTorch/tree/048cb8c23a4dacadffd3db53f25aa7f4660087cf

GitHub - Lornatang/YOLOv4-PyTorch: Pytorch implements yolov4.Good performance, easy to use, fast speed.

 

yolov4 는 tensorflow로 작성되었다.

다른 사람이 pytorch용으로 수정한 버전을 가지고 온다.

 

git clone https://github.com/Lornatang/YOLOv4-PyTorch/tree/048cb8c23a4dacadffd3db53f25aa7f4660087cf

 

YOLOv4-PyTorch / configs / COCO-Detection / yolov4.yaml 파일을 복사해서 yolov5 / models에 붙여넣기한다.

yolov4.yaml 파일을 아래와 같이 수정해준다.

 

number_classes를 nc로 수정하는것만 해주면 된다.

 

# 2번줄
number_classes: 80  >> nc: 80

# 61번줄
[[30, 33, 36], 1, Detect, [number_classes, anchors]],  # Detect(P3, P4, P5)
  >> [[30, 33, 36], 1, Detect, [nc, anchors]],

 

 

2. common.py 수정

yolov4에 사용하는 함수는 yolov5에서 사용하지 않는다.

따라서 yolov5에 있는 common.py에 yolov4의 함수를 추가해줘야한다.

다행히 큰 변화는 없다.

 

common.py 경로 : yolov5/models

 

# BottleneckCSP is for yolov4 
class BottleneckCSP(nn.Module):
    # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
        super(BottleneckCSP, self).__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
        self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
        self.cv4 = Conv(2 * c_, c2, 1, 1)
        self.bn = nn.BatchNorm2d(2 * c_)  # applied to cat(cv2, cv3)
        self.act = nn.LeakyReLU(0.1, inplace=True)
        self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])

    def forward(self, x):
        y1 = self.cv3(self.m(self.cv1(x)))
        y2 = self.cv2(x)
        return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))

 

혹시 이 포스팅에 있는 깃허브가 아니라 다른 깃허브에서 yolov4를 끌어왔다면 파일 내용이 좀 다를 수 있다.

필요한 부분들은 찾아서 가지고 오면 된다.

 

class BottleneckCSP() 함수는 yolov4/yolov4_pytorch/model/module/neck.py 에서 찾아서 가지고 왔다.

 

 

3. Train Val Test Inference

이후 내용은 이전에 포스팅 했던것과 동일하게 진행하면 된다.

다만 yolov4를 yolov5 위치에서 실행하는 것이기 때문에 필요한 파라미터가 조금 다르다.

 

이전 포스팅 참고 : [object detection] 이미지 라벨링, 폴리곤 작업후 (CVAT, json) 시각화(yolov5, convert2Yolo) 실습 코드

 

python train.py --img 이미지사이즈 \
--batch 배치사이즈 \
--epochs 에폭수 \
--data custom.yaml파일 경로 \
--cfg 사용할 모델의 yaml파일 경로 \
--weights 전이 학습에서 사용할 wtight \
--name 학습된 정보를 runs 폴더안에 저장할 이름 \
--project wanbd에 저장할 프로젝트명

 

파라미터들이 많이 있는데 yolov4 같은 경우는 모델을 가져온게 아니고 yaml 파일만 가져왔기 때문에 --cfg를 통해 yolov4.yaml 를 설정해주고 --weights는 빈 값 ' '으로 둔다. yolov5는 --cfg는 필요없고 --weights를 통해 모델를 직접 설정해준다.

 

# yolov4
python train.py --data ./custom.yaml --cfg ./models/yolov4.yaml --weights '' --name plastic_v4

# yolov5m
python train.py --data ./custom.yaml --weights yolov5m.pt --name plastic_v5m

 

이 둘의 차이점은 yolov4는 이전에 학습된 것 없이 아주 처음 과정부터 학습하는 상황이라면 yolov5는 기존 학습된 모델(pre_trained)을 불러와서 내 데이터를 학습시키는 것이라는 점이다.

 

yolov4 학습 결과

yolov5 학습 결과

 

위와 같이 결과 이미지에 나오는 loss를 보면 학습할 때 yolov5가 좀 더 매끄러운 곡설을 그리며 잘 학습되었다는 것을 알 수 있다.

mAP도 yolov5는 일괄적으로 0.995를 보이며 좋은 학습을 보였다.