register 기법 적용된 ViT 사용시 유의 사항

import torch
import timm

if __name__ == '__main__':
    dinov2_model_wo_reg = timm.create_model('vit_base_patch14_dinov2.lvd142m', pretrained=True)
    dinov2_model_w_reg = timm.create_model('vit_base_patch14_reg4_dinov2.lvd142m', pretrained=True)

    input = torch.randn(1, 3, 518, 518)
    output_wo_reg = dinov2_model_wo_reg.forward_features(input)
    output_w_reg = dinov2_model_w_reg.forward_features(input)

    print(output_wo_reg.shape)
    print(output_w_reg.shape)

 

위 코드에 대한 출력 결과는 아래와 같다.

 

torch.Size([1, 1370, 768])
torch.Size([1, 1374, 768])

 

vit_base_patch14_reg4_dinov2.lvd142m 모델은 register 기법에 사용되는 4개의 토큰이 추가로 붙기에, 반환되는 토큰의 수가 그렇지 않은 모델보다 4개 더 많다.

 

이 토큰을 사용할지 말지도 좀 중요할 거 같아서 블로그에 따로 기록해둔다.

 

논문에서는 버리라고 나온다.

 

버릴려면 어떻게 해야할까

 

 

논문 그림에서는 cls token이랑 reg token이 뒤에 있지만 실제로 구현은 앞단에 concat하는 식으로 구현돼있다.

 

outputs[:, 5:] 이런식으로 인덱싱해야 cls token, reg token 다 날려버릴수가 있다. 이것도 실제 사용 케이스에 따라 ablation study를 모조리 진행해야할 거 같다.