[Penn-Fudan] Object detection 학습하기 - 2
이전 포스팅 내용이 궁금하시다면 아래의 링크 참고하시면 감사하겠습니다 !
- [Penn-Fudan] Object detection 학습하기 - 1
저번 포스팅에 이어 이번 포스팅은 Pen-Fudan 데이터셋을 학습할 모델을 정의하고 코드로 작성해볼 것이다.
모델은 Mask R-CNN을 사용하되, 2가지 방법으로 구현해 볼 것이다.
- Pretrain된 Mask R-CNN 모델을 불러온 뒤, head를 Faster R-CNN으로 변경하고 마지막 layer만 학습
- Backbone을 MobileNet v2로 변경한 뒤 학습
2. 학습 모델 정의
2.1 Pretrain된 모델을 불러온 뒤, head를 변경 후 마지막 Layer만 학습
해당 코드는 다음과 같다.
import torchvision
import torch.nn as nn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
class PreMaskRCNN(nn.Module):
def __init__(self, num_classes=2):
super().__init__()
# load a Pre-trained Mask R-CNN model on COCO dataset
self.model = torchvision.models.detection.maskrcnn_resnet50_fpn(weights="DEFAULT")
# For replacing the classifier (last layer) => 2 (background + pedestrian)
self.num_classes = num_classes
# get number of input features for the classifier
in_features = self.model.roi_heads.box_predictor.cls_score.in_features
# Replace the pre-trained head with new one
self.model.roi_heads.box_predictor = FastRCNNPredictor(in_features, self.num_classes)
in_features_mask = self.model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
# and replace the mask predictor with a new one
self.model.roi_heads.mask_predictor = MaskRCNNPredictor(
in_features_mask,
hidden_layer,
self.num_classes
)
def forward(self, imgs, targets=None):
if targets:
return self.model(imgs, targets)
else:
return self.model(imgs)
먼저, COCO 데이터셋으로 pretrain된 Mask R-CNN 모델을 불러와준다.
여기서 resnet50은 backbone을 resnet 50을 사용했다는 것을 의미하고, fpn은 head로 FPN(Featuer Pyramid
Network)를 사용했다는 것이다.
그리고 head를 Fast R-CNN의 box predictor로 바꿔준다.
그 후, 256개의 hidden layer를 추가하고 수정된 head와 classifer 파라미터를 업데이트해준다.
2.2 Backbone을 MobileNet V2로 변경하여 학습
이번에는 backbone을 바꿔서 학습할 것이다. 코드는 아래와 같다.
import torchvision
import torch.nn as nn
from torchvision.models.detection import MaskRCNN
from torchvision.models.detection.anchor_utils import AnchorGenerator
class MaskRCNN_MobileNet_v2_bn(nn.Module):
def __init__(self, num_classes=2):
super().__init__()
self.num_classes = num_classes
# Load pretrained mobilenet_v2 and extract features
backbone = torchvision.models.mobilenet_v2(weights="DEFAULT").features
backbone.out_channels = 1280
# Make RPN generate 5 x 3
anchor_generator = AnchorGenerator(
sizes=((32, 64, 128, 256, 512),),
aspect_ratios=((0.5, 1.0, 2.0),)
)
# Define the feature maps that we will use to perform the roi cropping
roi_pooler = torchvision.ops.MultiScaleRoIAlign(
featmap_names=['0'],
output_size=7,
sampling_ratio=2
)
# Put the pieces together
self.model = MaskRCNN(
backbone,
num_classes=self.num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler
)
def forward(self, imgs, targets=None):
if targets:
return self.model(imgs, targets)
else:
return self.model(imgs)
여기서 우리는 빈 Mask R-CNN의 구조를 가져와서 필요한 모듈을 우리가 넣어줄 것이다. 우리가 채워야할 모듈은
"backbone", "rpn_anchor_generator" 그리고 "box_roi_pool"이다.
각 값들에 대한 설명을 간략히 적어보면 다음과 같다.
- backbone : 입력된 이미지로부터 고차원 feature를 추출. 위의 코드에서는 MobileNet v2를 사용
- rpn_anchor_generator : RPN(Region Proposal Network)에 사용할 Anchor 파라미터를 정의.
위의 코드에서는 총 5개의 Size와 3개의 비율값을 통해 object detection을 진행 - box_roi_pool : 객체의 bounding box를 정밀하게 추정하기 위해, feature map에서 RoI를 pooling하는
방식을 정의. 위 코드에서는 MultiScaleAlign을 통해 다양한 크기의 featuremap에서 효과적으로 RoI를 pooling
나는 이 두 개의 모델 코드를 한 개의 Python 파일에 작성하여 저장해줬다.
이제 정의된 두 모델에 대해 살펴보자.
먼저, 첫 번째 방법에서 구현된 모델의 정보를 보자.
위의 사진에서 알 수 있듯 모델의 학습 가능한 파라미터 수는 대략 43.7M(Milion) 정도다. 모델의 크기는 3794.11 MB다.
반면, 두 번째는 87.3M 정도의 학습 가능한 파라미터 수를 갖고 있지만, 모델의 크기는 3760.88 MB로 앞선 경우와
비슷하다.
데이터양의 충분하다면, MobileNet v2 bakcbone을 사용한 모델이 학습은 느리더라도 더 높은 정확도를 보여주겠지만,
데이터셋이 적어서 결과가 어떨지 모르겠다. 실제 예제에서도 데이터가 너무 적어서 첫 번째 모델로 진행했다고
적혀있긴하다.
그래도 pretrain된 backbone을 사용하는데 두 번째 방식도 잘되지 않을까 생각이 들었다.
이번 포스팅에서는 Mask R-CNN을 정의하는 2가지 방법에 대해서 다뤄보았다. 다음 포스팅에서는 정의된 모델을 학습하는
코드를 다루겠다.
내용에 대한 질문이나 오류에 대해 답글 달아주시면 감사하겠습니다! :)