DropBlock, 당신은 최근 Vision 모델에서 Dropout을 본 적이 있는가

✅ Intro

Computer Vision 관련 아키텍처가 발전함에 따라 Overfitting을 막기 위해 Batch Normalization, 혹은 Layer Normalization을 사용하는 추세입니다. 하지만, 이전에는 Dropout이 있었죠. Dropout은 어느샌가부터 마지막 fully-connected layer를 제외하고는 Computer Vision에서 잊혀 갔습니다.

이러한 이유에는 기존 Dropout은 Feature Map의 공간적인 특성을 고려하지 않는 부분에 있습니다. 픽셀 별로 랜덤 하게 Drop 시키는 경우를 생각해 보면 이해가 됩니다. Feature Map에서 고작 한 픽셀을 Drop 시킨다는 건 의미가 없을 수 있습니다.

하지만, 그림(c)과 같이 Feature Map의 특성을 고려하여 Drop을 한다면 효과가 있지 않을까요?

위와 같은 방법을 제안한 연구(논문)에서는 이 Method를 토대로 ResNet을 활용하여 실험을 진행했습니다.

결과는 위와 같이 Drop을 하는 Feature Map(Block size)가 커질수록 성능이 좋아짐을 보였습니다.

연구에서는 이러한 Method를 DropBlock이라 하고, 이 포스팅에서는 DropBlock을 구현해 보겠습니다.

---

✅ Implementation

DropBlock을 처음에 구현할 때는 너무 막연하게 Tensor를 output으로 돌려주는 콘셉트에 기반하여 class를 만들어주었습니다. 하지만, 이 방식은 너무 오래 걸릴뿐더러 'DropBlock에 학습시킬 요소가 없는데 왜 Module 형태의 class를 만들어준 걸까?'라는 결론을 도출했습니다. 이러한 생각은 앞으로 아키텍처를 구현, 혹은 필요한 구성요소가 있을 때도 고려해야 할 요소로 자리 잡았습니다. (class로 구현한 부분은 포스팅 하단에 첨부하겠습니다.)

1. Drop 시킬 Center Pixel 선정

우선 블록 단위로 Drop을 시키기 위해서는 그 중심이 되는 Center Pixel을 선정해야 합니다. 선정하는 방식은 Dropout과 같으며, 다른 점은 Feature Map의 크기를 고려하여 각 픽셀이 Drop 될 확률을 구해줘야 합니다. 그래서, Keep Probability에 대해 각 픽셀이 Drop 될 확률(Gamma)을 구하는 공식은 아래와 같습니다.

$$ \gamma = \frac{1-keep\_prob}{block\_size^2}\frac{feat\_size^2}{(feat\_size-block\_size+1)^2} $$

위 식에서 알 수 있듯이 feat_size라는 input 상수를 제외한 keep_prob, block_size는 모두 사용자가 지정할 Hyperparameter에 속합니다.

gamma = (1.0-keep_prob)*H*W / ((block_size**2) * ((H - block_size + 1) * (W - block_size + 1)))

2. 베르누이 분포를 통해 0(Drop) or 1(Keep)

여기서 베르누이 분포를 더 자세히 다루기보다는 베르누이 분포라는 것은 Gamma 값을 input으로 넣었을 때, 랜덤 한 확률에 따라 0 or 1의 output을 주는 일종의 도구라 해석하는 것이 좋습니다.

noise = torch.empty((N, C, H - block_size + 1, W - block_size + 1), dtype=input.dtype, device=input.device)
noise.bernoulli_(gamma)

3. Padding & Getting Block

이를 통해 Drop이 될 각 Center Pixel을 선정하였으나, Block Size에 따라 Feature Map의 범위에 벗어나지 않도록 해주는 Padding 작업이 필요합니다.

또한, MaxPooling을 사용해 일종의 트릭과 유사하게 Center Pixel에서 Block으로 확장시킵니다.

noise = F.pad(noise, [block_size // 2] * 4, value=0)
noise = F.max_pool2d(noise, stride=(1, 1), kernel_size=(block_size, block_size), padding=block_size // 2)

4. Output

Drop 되는 Block의 Pixel이 positive로 잡혀있기 때문에 원본 이미지에 대해 이를 Reverse 시켜서 Drop이 되도록 곱합니다.

noise = 1 - noise
return input * noise

 

위 과정을 통해 나오는 함수는 아래와 같이 생겼습니다.

def dropBlock(input, block_size, keep_prob=0.9):
    N, C, H, W = input.size()
    block_size = min(block_size, W, H)
    gamma = (1.0-keep_prob)*H*W / ((block_size**2) * ((H - block_size + 1) * (W - block_size + 1)))
    noise = torch.empty((N, C, H - block_size + 1, W - block_size + 1), dtype=input.dtype, device=input.device)
    noise.bernoulli_(gamma)

    noise = F.pad(noise, [block_size // 2] * 4, value=0)
    noise = F.max_pool2d(noise, stride=(1, 1), kernel_size=(block_size, block_size), padding=block_size // 2)
    noise = 1 - noise
    return input * noise

 

DropBlock 함수는 제가 직접 구성한 것이 아니라 PyTorch Documentation에서 제공하는 DropBlock2D 함수를 제가 사용할 수 있도록 커스터마이징 한 것입니다. 그래서, 이를 통해 MaxPooling, Pad와 같은 Functional을 '트릭의 용도로 사용할 수 있겠구나'를 알아갑니다.

---

✅ Use DropBlock Function

랜덤으로 픽셀을 잡아 생성된 이미지에 대해 DropBlock을 사용하면 아래와 같습니다. 논문에서는 '각 채널 별로 독립적인 DropBlock을 갖게 하는 것이 더 좋은 성능을 보였다'라고 하였기 때문에 위 함수도 채널 별로 독립적으로 사용하도록 합니다.

channels=3
sample = torch.rand((3, channels, 224, 224))
sample_db = dropBlock(sample, block_size=51, keep_prob=0.9)

ch_cnt = 1
for i in range(1, channels+1):
    plt.subplot(channels, 2, i*2-1)
    plt.imshow(to_pil_image(sample[0][i-1]))
    plt.axis('off')
    plt.title(f'ch{ch_cnt} orig')

    plt.subplot(channels, 2, i*2)
    plt.imshow(to_pil_image(sample_db[0][i-1]))
    plt.axis('off')
    plt.title(f'ch{ch_cnt} DropBlock')
    ch_cnt += 1

plt.show()

---

✅ Class version

Function 버전 이전에 구현했던 Class 버전에 대해 작동은 하지만, Training을 할 경우 엄청나게 오랜 시간이 걸립니다. Tensor를 사용하고, 해당 Tensor를 Device로 넘기는 과정을 학습 중에 하도록 설정해서 그랬습니다.

(그래도 일단 올려둡니다..ㅎ)

class DropBlock(nn.Module):
    def __init__(self, block_size, keep_prob=0.9, sync_channel=False):
        super(DropBlock, self).__init__()
        '''
        block_size : feature에서 drop 시킬 block의 크기, 반드시 홀수여야 함.
        keep_prob : 계속 activation 시킬 Probability
        논문에서는 Keep Probability에 대해 학습을 하면서 1부터 알맞는 값까지
        선형적으로 학습하여 적합한 p값을 찾아야 한다 하지만, 그 부분은 구현이 꽤
        어려워서 0.9를 default로 모델링을 한다.
        '''
        self.block_size = block_size
        self.keep_prob = keep_prob
        self.sync_channel = sync_channel
        self.padding_size = (self.block_size-1)//2

    def getGamma(self, feat_size):
        '''
        Gamma의 역할
        Traditional한 Dropout에서 1-keep_prob을 통해 베르누이 분포를 계산하여
        Dropout을 진행하는데 이 부분이 DropBlock에서는 영역적으로
        block_size에 따라 Drop하기 때문에 이 부분을 고려하여 아래와 같이
        Gamma를 계산하여 베르누이 분포의 Drop 시킬 확률 값으로 넘기도록 한다.
        '''
        return (1.0-self.keep_prob)/(self.block_size**2)*(feat_size**2)/((feat_size-self.block_size+1)**2)

    def dropMask(self, feat_size):
        '''
        여기서 (1-keep_prob)기반 Gamma를 토대로 Bernoulli를 쓰는 것이기 때문에
        1이 Drop할 Center Pixel이다.
        '''
        mask = torch.distributions.Bernoulli(probs=self.getGamma(feat_size)).sample((feat_size, feat_size))
        return mask

    def outOfRegion(self, mask_pixel):
        '''
        Mask할 Region들이 Block Size로 인해 feature map을 넘어가지 않도록
        feature map 내에서 fully하게 Drop할 수 있도록 테두리 부분에
        Mask 픽셀이 있는 경우 이를 제거 한다.
        '''
        mask_pixel[0:self.padding_size, :] = 0.
        mask_pixel[-self.padding_size:, :] = 0.
        mask_pixel[:, 0:self.padding_size] = 0.
        mask_pixel[:, -self.padding_size:] = 0.
        return mask_pixel

    def getRegion(self, mask_region):
        masking_li = []
        for i in range(0, mask_region.shape[0]):
            for j in range(0, mask_region.shape[1]):
                if mask_region[i][j] == 1.0:
                    masking_li.append((i, j))

        for (i, j) in masking_li:
            mask_region[i-self.padding_size:i+self.padding_size+1,j-self.padding_size:j+self.padding_size+1]=1.0
        return mask_region
        
    def forward(self, x):
        feat_size = x.shape[-1]
        n_channels = x.shape[-3]
        
        if self.sync_channel:
            '''
            논문 실험에서는 channel 다 통합한 같은 Masking보다
            독립적으로 하는 게 더 성능이 좋다해서 이 부분은 사용 안 할 듯
            '''
            mask = torch.where(\
                    self.getRegion(\
                        self.outOfRegion(\
                            self.dropMask(feat_size))) == 1, 0, 1).float()
            x = x * mask
            return x
        else:
            # Channel에 따라 독립적으로 Dropout
            mask = torch.stack(
                [torch.where(\
                    self.getRegion(\
                        self.outOfRegion(\
                            self.dropMask(feat_size))) == 1, 0, 1).float()\
                                 for _ in range(n_channels)], dim=0)
            x = x * mask
            return x

---

✅ Reference

[1] DropBlock: A regularization method for convolutional networks

[2] DropBlock2d in PyTorch