Pytorch로 구현한 Fast R-CNN 모델

이번 포스팅에서는 gary1346aa님의 github repository에 올라온 pytorch로 구현한 Fast R-CNN 코드를 분석해보도록 하겠습니다. jupyter notebook으로 작성되어 있어 코드를 상대적으로 읽기가 편했던 것 같습니다. 하지만 학습과 detection을 위한 데이터셋을 제공하지 않아, 일부 코드를 이해하기가 어려웠던 것 같습니다. 그럼에도 Fast R-CNN 모델의 핵심 아이디어가 코드로 잘 구현되어 있어 새로 배운 내용들이 많았습니다. 전체 코드를 살펴보기보다는 Fast R-CNN 모델의 핵심 아이디어 위주로 코드를 분석해보도록 하겠습니다. Fast R-CNN 모델에 대한 자세한 설명은 Fast R-CNN 논문 리뷰 포스팅을 참고하시기 바랍니다. 

---

1) RoI(Region of Interest) pooling

Fast R-CNN 

먼저 Fast R-CNN 모델의 가장 핵심 아이디어라고 할 수 있는 RoI(Region of Interest) pooling에 대해 살펴보도록 하겠습니다. 구체적으로 코드를 살펴보기에 앞서 전체 네트워크에서 RoI pooling을 적용하는 시점을 짚고 넘어가도록 하겠습니다. 이전 포스팅에서 살펴보았듯이, Fast R-CNN 모델은 원본 이미지를 pre-trained된 VGG16 모델에 입력합니다. 이 때 마지막 max pooling layer를 RoI pooling layer로 대체하여 고정된 크기의 feature map을 다음 fc layer에 전달합니다. RoI pooling을 수행하는 feature map의 크기는 14x14x512입니다. 이 점을 유념하고 코드를 살펴보면 이해하기가 더 수월할 것 같습니다. 

class SlowROIPool(nn.Module):    
    def __init__(self, output_size):
        super().__init__()
        self.maxpool = nn.AdaptiveMaxPool2d(output_size)
        self.size = output_size
    
    # images : 원본 이미지
    # rois : region of interests
    # roi_idx : 
    def forward(self, images, rois, roi_idx):
        n = rois.shape[0] # region of interest의 수
        
        # 고정된 크기로 들어오기 때문에 전부 다 14x14
        h = images.size(2) # h : feature map height
        w = images.size(3) # w : feature map width
        
        # region of interst의 (x1, y1, x2, y2)의 행렬
        # 상대 좌표로 들어옴
        x1 = rois[:,0]
        y1 = rois[:,1]
        x2 = rois[:,2]
        y2 = rois[:,3]
        
        # region of interest의 상대좌표를 feature map에 맞게 절대좌표로 변환함
        x1 = np.floor(x1 * w).astype(int)
        x2 = np.ceil(x2 * w).astype(int)
        y1 = np.floor(y1 * h).astype(int)
        y2 = np.ceil(y2 * h).astype(int)

RoI pooling을 수행하는 SlowROIPool 클래스의 forward 메서드에서 각 파라미터는 다음과 같습니다. 

 

• images : 14x14 크기의 512개의 feature map 리스트
• rois : region of interest(=region proposals)의 상대좌표
• roi_idx : region of interest의 index 리스트

저희가 살펴보는 코드의 데이터셋은 미리 Selective search 알고리즘을 이미지에 적용하여 region of interest를 pkl 형식으로 제공하고 있습니다.  이 때 region of interest의 좌표는 원본 이미지 크기에서 region of interest가 차지하는 비율 형식으로 저장되어 있습니다. 따라서 rois에 저장된 값들은 0~1 사이의 값을 가집니다. 이를 feature map(=14x14)의 크기에 맞게 feature map의 width, height를 곱해주고, numpy에서 제공하는 floor(내림), ceil(올림) 함수를 활용하여 feature map 내 region proposal이 encode하는 영역을 찾습니다.

class SlowROIPool(nn.Module):    
    def __init__(self, output_size):
        (...)
    def forward(self, images, rois, roi_idx):
    	(...)
		res = []
        
        # region of interest의 수만큼 순회
        for i in range(n):
            img = images[roi_idx[i]].unsqueeze(0) # roi_idx i번째 해당하는 feature map
            img = img[:, :, y1[i]:y2[i], x1[i]:x2[i]] # 잘라내기
            img = self.maxpool(img) # adaptive average pooling
            res.append(img)
        res = torch.cat(res, dim=0)
        return res # 7x7x(# of region proposals)

그 다음 전체 region of interest을 feature map에 대하여 RoI Projection을 수행하여 관심 영역을 추출하고 max pooling을 수행합니다. 이 때 일반적인 max pooling을 수행하는 것이 아니라, input feature map의 크기와 output feature map의 크기를 고려하여 stride와 kernel의 크기를 조정하는 Adaptive max pooling을 수행합니다. 이를 통해 고정된 크기(7x7)의 feature map을 출력하게 됩니다. 이러한 과정을 거쳐 최종적으로 7x7 크기의 feature map이 region of interest의 수만큼 저장된 리스트가 반환됩니다. 

2) Initializing pre-trained network

다음으로 pre-trained된 VGG16 모델을 load한 후 detection task에 맞게 네트워크를 수정하는 코드를 살펴보도록 하겠습니다. 

class RCNN(nn.Module):
    def __init__(self):
        super().__init__()

        rawnet = torchvision.models.vgg16_bn(pretrained=True)  # pre-trained된 vgg16_bn 모델 가져오기 
        self.seq = nn.Sequential(*list(rawnet.features.children())[:-1]) # 마지막 max pooling 제거
        self.roipool = SlowROIPool(output_size=(7, 7)) # 마지막 pooling layer, roi pooling으로 대체
        self.feature = nn.Sequential(*list(rawnet.classifier.children())[:-1])  # 마지막 fc layer 제거

        _x = Variable(torch.Tensor(1, 3, 224, 224))
        _r = np.array([[0., 0., 1., 1.]])
        _ri = np.array([0])
        _x = self.feature(self.roipool(self.seq(_x), _r, _ri).view(1, -1)) # 7x7x(# of region proposals)
        
        feature_dim = _x.size(1) 
        self.cls_score = nn.Linear(feature_dim, N_CLASS+1) # classifier
        self.bbox = nn.Linear(feature_dim, 4*(N_CLASS+1)) # bounding box regressor
        
        self.cel = nn.CrossEntropyLoss()
        self.sl1 = nn.SmoothL1Loss()

    def forward(self, inp, rois, ridx):
        res = inp # images
        res = self.seq(res) # ~pre-pooling
        res = self.roipool(res, rois, ridx) # roi pooling
        res = res.detach() # 연산 x
        res = res.view(res.size(0), -1)
        feat = self.feature(res) # fc layers

        cls_score = self.cls_score(feat) # classification result
        bbox = self.bbox(feat).view(-1, N_CLASS+1, 4) # bounding box regressor result
        
        return cls_score, bbox

• torchvision 라이브러리에서 제공하는 pre-trained된 VGG16 모델을 load해줍니다. 코드 원작자 분은 batch normalization이 추가된 VGG16 모델을 load했습니다.
• 네트워크에서 마지막 pooling layer를 제거한 후 위에서 정의한 RoI pooling layer로 대체해줍니다. 이 때 output feature map의 크기를 7x7이 되도록 설정합니다.  
• 마지막 fc layer를 제거합니다. 이는 Classifier와 Bounding box regressor를 추가하기 위함입니다. 
• Variable을 통해 VGG 모델에 입력되는 데이터의 크기(224x224 크기의 RGB 이미지)를 정의합니다. 원본 데이터를 conv layer, roi pooling layer, fc layer에 순차적으로 입력하여 고정된 크기(7x7)의 feature vector를 얻습니다. 
• feature vector를 Classifier와 Bounding box regressor에 입력합니다. 
• Classifier와 Bounding box regressor의 loss function을 각각 정의(Crossentropy, Smooth L1)합니다. 

3) Multi-task loss

class RCNN(nn.Module):
    def __init__(self):
        (...)

    def forward(self, inp, rois, ridx):
        (...)

    def calc_loss(self, probs, bbox, labels, gt_bbox):
        loss_sc = self.cel(probs, labels) # crossentropy loss
        
        lbl = labels.view(-1, 1, 1).expand(labels.size(0), 1, 4)
        mask = (labels != 0).float().view(-1, 1).expand(labels.size(0), 4)
        loss_loc = self.sl1(bbox.gather(1, lbl).squeeze(1) * mask, gt_bbox * mask)
        
        # multi-task loss, crossentropy loss, smooth l1 loss
        lmb = 1.0
        loss = loss_sc + lmb * loss_loc
        
        return loss, loss_sc, loss_loc

 

다음으로 Multi-task loss를 구현한 부분을 살펴보도록 하겠습니다. Classifier의 loss는 위에서 정의한 Crossentropy loss를 통해 구할 수 있습니다. 하지만 Bounding box regressor의 loss를 구하는 부분은 살짝 복잡합니다. Multi-task loss를 다시 한 번 살펴보도록 하겠습니다. 

 

$$L(p, u, t^u, v) = L_{cls}(p, u) + \lambda[u \ge 1]L_{loc}(t^u, v)$$

 

Bounding box regressor의 loss에는 학습 데이터가 positive/negative sample 여부를 알려주는 index paramter $u$를 곱해져 있습니다. $u$는 구현된 코드에서 mask 변수에 해당합니다. labels 변수에는 학습 데이터의 positive/negative sample 여부를 알려주는 배열이 저장되어 있습니다. mask 변수는 labels에 저장된 배열을 bounding box에 대한 정보가 저장된 배열의 크기에 맞게 변환한 배열입니다. mask를 예측 bounding box에 해당하는 bbox 변수와, ground truth box에 해당하는 gt_bbox 변수에 곱해준 후 Smooth L1 loss를 구합니다. 데이터셋이 제공되지 않아 입출력 데이터가 명확하지 않아, 이 부분이 특히 이해하기 어려웠던 것 같습니다😭. 

 

다음으로 두 loss 사이의 가중치를 조정하는 lamba 파라미터에 해당하는 lmb 변수를 1로 설정한 후, 앞에서 구한 두 loss를 더해 최종적으로 multi-task loss를 반환합니다. 

4) Train Fast R-CNN 

def train_batch(img, rois, ridx, gt_cls, gt_tbbox, is_val=False):
    sc, r_bbox = rcnn(img, rois, ridx) # class score, bbox
    loss, loss_sc, loss_loc = rcnn.calc_loss(sc, r_bbox, gt_cls, gt_tbbox) # losses
    fl = loss.data.cpu().numpy()[0]
    fl_sc = loss_sc.data.cpu().numpy()[0]
    fl_loc = loss_loc.data.cpu().numpy()[0]

    if not is_val:
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
    return fl, fl_sc, fl_loc

다음으로 각 batch별로 Fast R-CNN 모델이 학습하는 과정을 구현한 코드를 살펴보도록 하겠습니다. 학습 과정은 크게 어렵지 않습니다. 원본 이미지와 roi를 위에서 정의한 Fast R-CNN 모델에 입력한 후, loss를 구하는 과정이 구현되어 있습니다. 

def train_epoch(run_set, is_val=False):
    I = 2   # number of image : 2
    B = 64  # number of rois per image : 64
    POS = int(B * 0.25)  # positive samples : 16
    NEG = B - POS # negative samples : 48
    
    # shffle images
    Nimg = len(run_set)
    perm = np.random.permutation(Nimg)
    perm = run_set[perm]
    
    losses = []
    losses_sc = []
    losses_loc = []
    
    # 전체 이미지를 I(=2)개씩만큼 처리
    for i in trange(0, Nimg, I):
        lb = i
        rb = min(i+I, Nimg)
        torch_seg = torch.from_numpy(perm[lb:rb]) # read 2 images
        img = Variable(train_imgs[torch_seg], volatile=is_val).cuda()
        ridx = []
        glo_ids = []

 

Fast R-CNN 모델을 epoch별로 학습시키는 과정에서 이전 포스팅에서 살펴본 Hierarchical sampling을 통해 적절한 학습 데이터를 sampling하고 있습니다. 원본 이미지 중 2장을 sampling한 후, 각 이미지에서 64장의 region of interest를 sampling하고 있습니다. 이처럼 같은 이미지에서 sampling한 region of interest는 forward, backpropagation 시, 연산과 메모리를 공유할 수 있습니다. 전체 region of interest 중에서 positive sample은 25%, 나머지는 negative sample로 구성하도록 지정합니다. 

(...)
        for j in range(lb, rb):
            info = train_img_info[perm[j]]
            
            # roi의 positive, negative idx에 대한 리스트
            pos_idx = info['pos_idx']
            neg_idx = info['neg_idx']
            ids = []

            if len(pos_idx) > 0:
                ids.append(np.random.choice(pos_idx, size=POS))
            if len(neg_idx) > 0:
                ids.append(np.random.choice(neg_idx, size=NEG))
            if len(ids) == 0:
                continue
            ids = np.concatenate(ids, axis=0)
            
            # glo_ids : 두 이미지에 대한 positive, negative sample의 idx를 저장한 리스트
            glo_ids.append(ids)
            ridx += [j-lb] * ids.shape[0]

        if len(ridx) == 0:
            continue
        glo_ids = np.concatenate(glo_ids, axis=0)
        ridx = np.array(ridx)
        (...)

positive/negative sample의 index가 저장된 리스트에서 지정한 positive/negative 수에 맞게 sampling합니다. 그리고 이미지 2장에서 sampling한 region of interest를 glo_ids 변수에 저장하여 epoch당 학습 데이터로 사용합니다.  

(...)
	rois = train_roi[glo_ids]
        gt_cls = Variable(torch.from_numpy(train_cls[glo_ids]), volatile=is_val).cuda()
        gt_tbbox = Variable(torch.from_numpy(train_tbbox[glo_ids]), volatile=is_val).cuda()

        loss, loss_sc, loss_loc = train_batch(img, rois, ridx, gt_cls, gt_tbbox, is_val=is_val)
        losses.append(loss)
        losses_sc.append(loss_sc)
        losses_loc.append(loss_loc)

    avg_loss = np.mean(losses)
    avg_loss_sc = np.mean(losses_sc)
    avg_loss_loc = np.mean(losses_loc)
    print(f'Avg loss = {avg_loss:.4f}; loss_sc = {avg_loss_sc:.4f}, loss_loc = {avg_loss_loc:.4f}')
    
    return losses, losses_sc, losses_loc

Hierarchical sampling을 통해 구성한 학습 데이터를 Fast R-CNN 모델에 입력한 후 loss를 구합니다. 학습 과정 자체는 크게 어렵지 않았으나, 학습 데이터를 sampling하는 과정이 살짝 복잡했던 것 같습니다. 

★ Regressor to Bounding box

def reg_to_bbox(img_size, reg, box):
    img_width, img_height = img_size
    bbox_width = box[:,2] - box[:,0] + 1.0
    bbox_height = box[:,3] - box[:,1] + 1.0
    bbox_ctr_x = box[:,0] + 0.5 * bbox_width
    bbox_ctr_y = box[:,1] + 0.5 * bbox_height

    bbox_width = bbox_width[:,np.newaxis]
    bbox_height = bbox_height[:,np.newaxis]
    bbox_ctr_x = bbox_ctr_x[:,np.newaxis]
    bbox_ctr_y = bbox_ctr_y[:,np.newaxis]

    out_ctr_x = reg[:,:,0] * bbox_width + bbox_ctr_x
    out_ctr_y = reg[:,:,1] * bbox_height + bbox_ctr_y

    out_width = bbox_width * np.exp(reg[:,:,2])
    out_height = bbox_height * np.exp(reg[:,:,3])

    return np.array([
        np.maximum(0, out_ctr_x - 0.5 * out_width),
        np.maximum(0, out_ctr_y - 0.5 * out_height),
        np.minimum(img_width, out_ctr_x + 0.5 * out_width),
        np.minimum(img_height, out_ctr_y + 0.5 * out_height)
    ]).transpose([1, 2, 0])

마지막으로 제가 개인적으로 중요하다고 생각한 코드를 살펴보고 마무리하겠습니다. Bounding box regressor는 예측 bounding box의 좌표값을 변환해주는 Scale-invariant transformation을 학습합니다. 이는 곧 모델이 학습한 결과가 예측 bounding box의 좌표값이 아니라, transformation임을 의미합니다. 따라서 학습한 결과를 다시 bounding box 형태로 변환해주는 별도의 작업이 필요합니다. 위의 코드에서는 reg 파라미터를 통해 Bounding box regressor가 학습한 결과(=transformation)를 할당하고, 이를 활용하여 bounding box 좌표 형태(x, y, w, h)로 변환해주고 있습니다. 

---

지금까지 pytorch로 구현한 Fast R-CNN 모델의 코드를 분석해봤습니다. 개인적으로 RoI pooling을 구현하는 부분이 상당히 흥미로웠습니다. 그리고 코드를 살펴보면서 gather, expand 등 pytorch에서 배열의 형태를 변환해주는 함수를 더 공부해야 될 필요성을 느꼈습니다. 데이터셋이 제공되었다면 조금 아쉬웠고, kaggle에 올라온 데이터셋을 활용해서 코드를 조금 수정해볼 계획입니다.

Reference

gary1346aa님의 github repository

Fast R-CNN 논문 리뷰

Adaptive max pooling에 대한 설명

Fast R-CNN 논문 리뷰 포스팅