Perceiver IO: Optical Flow

지난 포스팅에 이어 Perceiver IO의 적용(optical flow)에 대해 코드 단위로 알아보도록 하겠다. Huggingface에서 구현한 라이브러리를 중심으로 설명을 한다. Perceiver IO의 기본개념에 대해서는 아래의 이전포스팅을 확인해주길 바란다. 

 

Perceiver IO 기본개념

https://oculus.tistory.com/59?category=500795 

Perceiver IO: a scalable, fully-attentional model that works on any modality

Perceiver IO: Text Classification

https://oculus.tistory.com/60?category=500795 

Perceiver IO: Text Classification

Perceiver IO: Image Classification

https://oculus.tistory.com/61?category=500795 

Perceiver IO: Image Classification

 

 

---

Perceiver for optical flow

Optical Flow를 추정하는 테스크는 같은 Scene에서 두장의 image가 주어지고, 첫번째 이미지에 있는 픽셀들에 대해 2D displacement를 추정하는 테스크이다. Perceiver를 이용하면 간단하게 optical flow를 추정할 수 있다. 이전과 마찬가지로 Transformer library에 PerceiverForOpticalFlow 에 구현되어있다. 코드는 다음과 같다.

 

from torch import nn
from transformers import PerceiverModel
from transformers.models.perceiver.modeling_perceiver import PerceiverImagePreprocessor, PerceiverOpticalFlowDecoder

class PerceiverForOpticalFlow(nn.Module):
    def __init__(self, config):
        super().__init__(config)

        fourier_position_encoding_kwargs_preprocessor = dict(
            num_bands=64,
            max_resolution=config.train_size,
            sine_only=False,
            concat_pos=True,
        )
        fourier_position_encoding_kwargs_decoder = dict(
            concat_pos=True, max_resolution=config.train_size, num_bands=64, sine_only=False
        )
        
        image_preprocessor = PerceiverImagePreprocessor(
            config,
            prep_type="patches",
            spatial_downsample=1,
            conv_after_patching=True,
            conv_after_patching_in_channels=54,
            temporal_downsample=2,
            position_encoding_type="fourier",
            # position_encoding_kwargs
            fourier_position_encoding_kwargs=fourier_position_encoding_kwargs_preprocessor,
        )
        
        self.perceiver = PerceiverModel(
            config,
            input_preprocessor=image_preprocessor,
            decoder=PerceiverOpticalFlowDecoder(
                config,
                num_channels=image_preprocessor.num_channels,
                output_image_shape=config.train_size,
                rescale_factor=100.0,
                use_query_residual=False,
                output_num_channels=2,
                position_encoding_type="fourier",
                fourier_position_encoding_kwargs=fourier_position_encoding_kwargs_decoder,
            ),
        )

위에서 확인할 수 있듯이, preprocessor로 PerceiverImagePreprocessor를 사용하였고 decoder로는 PerceiverOpticalFlowDecoder를 사용하였다. 각각의 2개의 frame에 대해서 각 pixel에 대해 3x3 patch를 추출하였고, 이는 각 pixel당 3x3x3 = 27 사이즈의 value를 갖게 된다. 저자는 학습 resolution을 (368,496)로 설정하였다. 즉 input은 (batch_size, 2, 27, 368, 496)의 shape을 갖게 된다. 

 

먼저 Preprocessor는 channel dimension으로 frame들을 concatenate 하여 (batch_size, 368, 496, 54)의 형태로 만든다. 왜 channel dimension으로 합치는지에 대해서는 논문에 자세히 설명되어 있다. 그다음, spatial dimension을 flatten 하여 (batch_size, 368*496, 54) = (batch_size, 182528, 54) 의 tensor를 얻는다. 여기에 258 dimension의 position embedding을 추가해주어 (batch_size, 182528, 322)를 만든다. 해당 텐서는 cross-attention을 수행할 때 사용된다. 

 

latent로는 (batch_size, 2048, 512) shape의 tensor을 사용해주었다. 마찬가지로 cross-attention을 수행해서 같은 사이즈의 latent를 만들어낸다. 그 후 24 self-attention layer(head = 16)를 사용해서 latent embedding을 update 해준다. 

 

이제 optical flow를 추정하기 위해서 latent의 final hidden state를 decode 해야 하는데, PerceiverOpticalFlowDecoder는 preprocess 된 input (batch_size, 182528, 322)을 decoder query로 사용한다. 이렇게 하면 Output으로 (batch_size, 182528, 322) 를 얻을 수 있고, 이를 다시 기존 이미지 사이즈 (batch_size, 368, 496, 2)로 바꾸어 준다. 

 

아래 예시들을 통해 결과를 확인해볼 수 있다. 

 

 

 

여기까지 Perceiver IO를 살펴보았다! 본 시리즈는 아래의 huggingface의 포스팅을 기반으로 작성되었습니다. 

 

https://huggingface.co/blog/perceiver

Perceiver IO: a scalable, fully-attentional model that works on any modality