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

1. 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용으로 수정한 버전을 가지고 온다.

 

      

      

      

    
    

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

 

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

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

 

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

 

      

      

      

    
    

      <python />
    
# 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. 2. common.py 수정

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

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

다행히 큰 변화는 없다.

 

common.py 경로 : yolov5/models

 

      

      

      

    
    

      <python />
    
# 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. 3. Train Val Test Inference

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

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

 

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

 

      

      

      

    
    

      <python />
    
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를 통해 모델를 직접 설정해준다.

 

      

      

      

    
    

      <python />
    
# 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를 보이며 좋은 학습을 보였다.