Perceiver IO: Image Classification

지난 포스팅들에 이어, 오늘은 perceiver IO를 Image Classification에 적용한 예제를 살펴보도록 하겠다. 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: Optical Flow

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

Perceiver IO: Optical Flow

 

---

Perceiver for Images

앞서 포스팅한 text classification과 다른 점은, input image를 임베딩으로 만들 때 사용되는 preprocessor 뿐이다! Perceiver 저자들은 image preprocessing에 대해서 3가지 방법을 시도하였다. 

 

• Pixel value를 flatten 하고, kernel size = 1의 convolution layer을 사용하며, learned absolute 1D position embedding을 사용
• Pixel value를 flatten 하고, fixed 2D Fourier position embedding을 사용
• 2D convolutional + maxpool에 2D Fourier position embedding을 사용

 

위의 3가지 방식 모두 Transformer library에 구현되어 있고, 아래 링크에서 확인해 볼 수 있다. 사용시에는 PerceiverImagePreprocessor의 config를 바꿔줘야 한다. 본 포스팅에서는 첫번째 방식을 적용한 코드를 설명한다. 

https://huggingface.co/docs/transformers/model_doc/perceiver

Perceiver

 

PerceiverForImageClassificationLearned 모델을 살펴보면서 자세히 들어가보자. 

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

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

        self.perceiver = PerceiverModel(
            config,
            input_preprocessor=PerceiverImagePreprocessor(
                config,
                prep_type="conv1x1",
                spatial_downsample=1,
                out_channels=256,
                position_encoding_type="trainable",
                concat_or_add_pos="concat",
                project_pos_dim=256,
                trainable_position_encoding_kwargs=dict(num_channels=256, index_dims=config.image_size ** 2),
            ),
            decoder=PerceiverClassificationDecoder(
                config,
                num_channels=config.d_latents,
                trainable_position_encoding_kwargs=dict(num_channels=config.d_latents, index_dims=1),
                use_query_residual=True,
            ),
        )

 

위의 코드에서 prep_type = "conv1x1"과 trainable position encoding으로 설정되어있는 것을 확인할 수 있다. 그래서 자세히 어떻게 동작이 되는 것일까? 이미지 배치를 모델에 넣어주는 경우를 생각해보자. 이때 input의 shape은 (batch, 3, 224, 224) 이다. 

 

PerceiverImagePreprocessor는 input에 1x1 convolution layer을 적용해서 (batch_size, 256, 224, 224) 의 텐서를 반환한다. Reshape을 통해 (batch_size, 224, 224, 256) 으로 바꿔주고 spatial dimension (W, H)을 flatten 시켜주면 (batch_size, 50176, 256) 의 텐서를 얻을 수 있다. 그 후, 1D positional embedding과 concatenate 시켜준다. 위에서 설정해준 바와 같이 position embedding dimension이 256이기 때문에, (batch_size, 50176, 512) shape의 텐서가 된다. 이제 해당 텐서는 latent와 cross attention을 수행한다. 

 

저자는 모든 image model에 대해서 (batch_size, 512, 1024) 의 latent를 사용하였다. 앞서 Text Classification에서 했던 것과 똑같이, cross-attention을 통해서 (batch_size, 512, 1024)의 output을 낸다. 그 후 6개의 self-attention block을 8번 거치며 (batch_size, 512, 1024)의 final hidden state을 얻게된다. 해당 output을 Classification logit으로 변경하기 위해 Text task와 동일하게 PerceiverClassificationDecoder을 사용하였다. latent에서 key+value를, (batch_size, 1, num_labels)의 trainable position embedding에서 query를 뽑는다. 그러면 (batch_size, 1, num_labels) 의 최종 결과를 얻게된다!

 

다음 포스팅에서는 Perceiver를 이용해서 optical flow를 추출하는 예제를 다뤄보도록 하겠다.