[DL/CV] NN interpolation의 backward propagation

Nearest Neighbor Upsampling

FPN(Feature Pyramids Network)에서는 top-down path를 통과할 때 high level feature와 low level feature를 섞기 위해서 upsampling 연산을 하는데, upsampling 방법으로는 Nearest Neighbor Upsampling을 사용한다.

Nearest Neighbor Upsampling은 다음과 같은 연산이다. (나중에 Upsampling을 한 번에 정리할 필요가 있다고 느낌)

이미지 출처: https://www.imageeprocessing.com/2017/11/nearest-neighbor-interpolation.html

간단하게 생각하면 upsampling 시, 채워야 하는 부분을 가까운 값의 copy로 채우는 연산이다.

---

NN(Nearest Neighbor) Upsampling의 Backward path

NN Upsampling (interpoplation) 은 deconvolution과 같은 learnable upsampling 기법이 아니다. 따라서 해당 레이어에는 어떠한 gradient가 계산될 만한 파라미터가 존재하지 않는다.

하지만, forward path 상에 있으므로 backward path 상에서도 당연히 gradient는 흐를 것이다. 하지만 정확히 어떤 식으로 흐르는가?에 대해 고민을 해보았을 때 답을 내리기 어려웠다. 따라서 1) 직관적으로 생각한 것, 2) 파이토치 내에서 구동하는 방식 을 글에 적어보고자 한다.

 

1) scatter and gather

일단 nearest neighbor interpolation을 수식으로 정리하면 다음과 같다. 화살표 이후만 참고하면 된다. 여기서 $x'$ 는  생성하고자 하는 값이고, $\lambda$는 각 점 $x_1$, $x_2$과 $x'$ 간의 거리를 의미한다. (여기서 $x_2 - x_1 = 1$으로 가정함)

출처: 이화여자대학교 open sw project 1강 강의자료

풀어서 설명하면, 가까운 위치에 있는 점의 값을 가져오겠다는 것이다 (1차원). 아래 그림을 보면 하나의 값이 4개로 흩뿌려지는 형태(scatter)임을 알 수 있다. 

 

여기서 잠시 max pooling 연산에 대한 forward-backward propagation을 보면,

이미지 출처: https://leonardoaraujosantos.gitbook.io/artificial-inteligence/machine_learning/deep_learning/pooling_layer

max pooling에서의 max node는 router로서 동작을 하고 있음을 알 수 있다. max값만 그 다음 path로 넘겨주기 때문에 backward path에서는 max값을 흘려보내준 노드로만 gradient가 흐르게 된다.

위의 max pooling의 사례를 보고 NN Upsampling의 node를 도식화해보면 다음과 같다.

즉 위에서 말했던 것처럼 forward path의 경우에는 scatter해주고, backward path로는 그 역연산으로 gather를 해준다. 하지만 이런 reduction의 경우에는 reduce를 어떻게 해줄지가 중요한데, 평균보다는 summation일 것이라는 생각을 했다 (scatter하면서 A의 영향을 흩뿌려주었는데 받아올 때는 평균을 해버린다면 gradient를 하나 더 조작한다?는 느낌이 들었다 → 논리적인 설명이 힘들다..)

따라서, gather되는 형식은 맞지만 어떤 식으로 gather가 되는가를 확실히 알아보기 위해 파이토치로 실험을 진행해보았다.

 

2) How it works in pytorch

실제로 간단한 test layer를 만들어주고, 파이토치 내에서는 leaf node가 아니고 intermediate variables에 대해서는 requires_grad = True이더라도 backward 시 gradient를 삭제해버린다. 따라서 retain_grad 옵션을 주어 해결하였다.

class Test(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(in_channels = 1, out_channels = 1, kernel_size = 1)
        self.upsample = torch.nn.UpsamplingNearest2d(scale_factor = 2)
        self.conv2 = torch.nn.Conv2d(in_channels = 1, out_channels = 1, kernel_size = 3)
        self.fc = torch.nn.Linear(4, 1)
    def forward(self,x):
        x = self.conv1(x)
        setattr(self, 'conv1_output', x)
        upsampled = self.upsample(x)
        setattr(self, 'upsampled', upsampled)
        output = self.conv2(upsampled)
        #print(output.shape)
        output = self.fc(output.flatten(1))

        return output

# experiments
x= torch.tensor([1,2,3,4,5,6,7,8], dtype = torch.float, requires_grad = False).view(2,1,2,2)
loss_fn = torch.nn.MSELoss()

test_model = Test()

output = test_model(x)
loss = loss_fn(output, torch.Tensor([0,0]).unsqueeze(1))

# feature map에 흐르는 grad를 출력할 수 있도록 retain_grad를 적용해줌 before backward()
test_model.conv1_output.retain_grad()
test_model.upsampled.retain_grad()

loss.backward()


f1_grad = test_model.conv1_output.grad # dl/df1
f2_grad = test_model.upsampled.grad # dl/d upsampled

 

위의 코드에 대한 수행 결과는 다음과 같다.

실제로 f2_grad, 즉 dl/d(upsampled)를 각각 2x2 sum pooling을 진행한 결과, f1_grad와 동일한 결과를 내는 걸 알 수 있었다.

이로써 파이토치에서는 nearest neighbor upsampling의 backward path를 실제 forward path 상에서 흘러들어왔던 원소로 gradient summation을 함을 알 수 있었다.

 

---

cf.) 위에서 언급한 "leaf node"와 "intermediate variable"은 무엇일까?

다음은 torchviz를 이용하여 위의 Test 클래스의 forward를 시각화한 것이다. 이를 확인해보면, AccumulateGrad가 적용되는 연산은 해당 연산이 그래프의 leaf node (더이상 backward 상에서 뻗어나가는 가지가 없음. 즉, forward 상에서 들어오는 가지가 없음) 임을 알 수 있다. 반면, 아까 내가 구하고자 했던 feature map (여기서는 UpsampleNearest2dBackward 정도..) 은 forward path 상 들어오는 가지와 뻗어나가는 게 존재한다. 즉 "intermediate variable"이다.