[VLM]PyTorch로 CLIP 구현

1. Image/Text Encoder로 Feature 추출하기

CLIP에는 Image/Text 서로 다른 두 종류의 입력이 들어온다. 따라서 각각을 처리하는 encoder도 따로 필요하다.

여기서는 Backbone model로 ViT_tiny와 BERT를 사용했다.

#ViT
vision_encoder = timm.create_model(
    "vit_tiny_patch16_224",
    pretrained=True
)

#BERT
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
text_encoder = AutoModel.from_pretrained("bert-base-uncased")

def get_image_feature(images):
    features = vision_encoder.forward_features(images)
    image_cls = features[:,0]
    return image_cls

def get_text_feature(texts):
    inputs = tokenizer(
        texts,
        padding = True,
        truncation = True,
        return_tensors="pt"
    )

    outputs = text_encoder(**inputs)
    text_cls = outputs.last_hidden_state[:,0]

    return text_cls
    
images = torch.randn(4, 3, 224, 224)

texts = [
    "a photo of a dog",
    "a photo of a cat",
    "a photo of a car",
    "a photo of a bird"
]

image_cls = get_image_feature(images)
text_cls = get_text_feature(texts)

ViT와 BERT를 각각 불러온 후 각각 이미지와 Text input을 집어넣어 feature 값을 추출한다. 일단 임의의 image vector와 text를 만들어 진행했다.

여기서 Image feature와 Text feature의 크기가 다르기 때문에 이대로 cosine similarity의 계산은 불가능하다. 따라서 Projection을 통해 같은 Embedding Space로 만들어줘야한다.

 

2. Projection으로 같은 Embedding Space 만들기

image_projection = nn.Linear(192,256)
text_projection = nn.Linear(768,256)

image_embedding = image_projection(image_cls)
text_embedding = text_projection(text_cls)

ViT_tiny를 거치고 온 데이터는 (B_size, 192), BERT를 거치고 온 데이터는 (B_size, 768)의 크기를 가지기 때문에 각각을 (B_size, 256)으로 만들어준다.

 

3. Normalize 후 [B, B] Similarity Matrix 만들기

이제 image embedding과 text embedding의 dimension이 같아졌기 때문에 두 vector를 비교할 수 있다.

CLIP에서는 similarity를 계산하기 전에 각 embedding을 normalize한다.

image_embedding = F.normalize(image_embedding, dim=-1)
text_embedding = F.normalize(text_embedding, dim=-1)

dim=-1은 Tensor의 마지막 dimension을 기준으로 normalize하겠다는 의미이다. 현재 [B, D]  shape를 갖고 있으므로 batch안의 모든 image와 모든 text를 cosine similarity로 서로 비교하면 CLIP에서 가장 중요하게 다뤘던 부분 중 하나인 [B ,B]의 Similarity Matrix가 생긴다.

logits = image_embedding @ text_embedding.T

여기까지 한 뒤 각각의 shape를 뽑아보면

print(image_embedding.shape)
print(text_embedding.shape)
print(logits.shape)

torch.Size([4, 256])
torch.Size([4, 256])
torch.Size([4, 4])

이렇게 제대로 진행 된 것을 확인 할 수 있다.

 

4. 양방향 Loss 계산하기

Similarity Matrix를 만들었다면 이제 이를 이용해 loss를 계산할 수 있다.

먼저 정답 label을 만든다.

labels = torch.arange(logits.size(0))

torch.arange(n)은 (0, 1, 2, ..., n-1) 형태의 Tensor를 만든다.

Batch size가 4라면,

labels
# tensor([0, 1, 2, 3])

가 된다. 이 값이 정답인 이유는 현재 batch가

Image0 ↔ Text0
Image1 ↔ Text1
Image2 ↔ Text2
Image3 ↔ Text3

이렇게 구성 돼 있기 때문에 similarity matrix의 대각선 위치가 모두 정답이라고 할 수 있다. 따라서 정답인 labels를 저렇게 만들어 주는 것이다.

#Image -> Text
loss_i = F.cross_entropy(logits, labels)
#Text -> Image
loss_t = F.cross_entropy(logits.T,labels)

loss = (loss_i + loss_t) / 2
print(loss)

마지막으로 만들어둔 labels와 logits를 이용해 loss를 각각 구한 뒤 평균을 해준다.

 

여기까지 코드([VLM]CLIP_step_by_step)

 

---

 

여기까지가 CLIP을 만들기 위한 기본적인 구조를 코드로 잡아보았고 각 코드가 어떤 역할을 하는지 확인 했으니 이를 이용해 CLIP 모듈을 만들어보자.

 

5. 하나의 CLIP Module로 만들기

먼저 선언부분부터 보면 

class CLIP(nn.Module):
    def __init__(self,vision_encoder,text_encoder,image_feature_dim=192,
        text_feature_dim=768, embed_dim = 256):
        super().__init__()

        self.vision_encoder = vision_encoder
        self.text_encoder = text_encoder

        self.image_projection = nn.Linear(image_feature_dim,embed_dim)
        self.text_projection = nn.Linear(text_feature_dim, embed_dim)

많이 필요하지 않다. Backbone으로 사용 할 모듈들을 저장하고 Image와 Text 각각의 Projection Layer를 만든다.

 

forward() 부분도 구현해보면

def forward(self, images, inputs):
        features = self.vision_encoder.forward_features(images)
        image_cls = features[:,0]

        outputs = self.text_encoder(**inputs)
        text_cls = outputs.last_hidden_state[:,0] 

        image_embedding = self.image_projection(image_cls)
        text_embedding = self.text_projection(text_cls)

        image_embedding = F.normalize(image_embedding, dim=-1)
        text_embedding = F.normalize(text_embedding, dim=-1)

        logits = image_embedding @ text_embedding.T

        return logits

여기서 inputs는 Text 데이터로 모듈 밖에서 tokenizer를 거친 뒤 들어오는 값을 의미한다. 

앞서 만들었던 것 처럼 Image와 Text 데이터를 각각의 encoder를 통해 feature로 만들어주고 projection을 통해 같은 차원으로 만들어준 뒤 normalization을 진행하고 logits을 만들어 낸다.

꽤나 간단하게 만들어지는 모듈이다.

 

이후 train loop를 만들어보면

import torch.optim as optim

images = torch.randn(4, 3, 224, 224)

texts = [
    "a photo of a dog",
    "a photo of a cat",
    "a photo of a car",
    "a photo of a bird"
]

inputs = tokenizer(
    texts,
    padding=True,
    truncation=True,
    return_tensors="pt"
    )

model = CLIP(vision_encoder,text_encoder)

# pretrained encoder freeze
for param in model.vision_encoder.parameters():
    param.requires_grad = False

for param in model.text_encoder.parameters():
    param.requires_grad = False

loss_fn = nn.CrossEntropyLoss()

optimizer = optim.Adam( filter(lambda p: p.requires_grad, model.parameters()),lr=1e-3)

num_epochs = 100 

model.train()

for epoch in range(num_epochs):
    logits = model(images, inputs)

    labels = torch.arange(logits.size(0))

    loss_i = loss_fn(logits,labels)
    loss_t = loss_fn(logits.T, labels)
    loss = (loss_i + loss_t) / 2

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

이제는 정말 많이 익숙해서 따로 설명하지 않아도 다 읽을 수 있을 것이라 생각한다. 

 

다만 여기서

for param in model.vision_encoder.parameters():
    param.requires_grad = False

for param in model.text_encoder.parameters():
    param.requires_grad = False
    
optimizer = optim.Adam( filter(lambda p: p.requires_grad, model.parameters()),lr=1e-3)

이것은 생소하게 느껴질 것이다.

 

우리가 pretrained backbone을 사용한다고 했을 때 이미 어느정도 좋은 성능을 가지고 있을 것이기 때문에 처음부터 다시 학습시키기 보다는

Image Backbone     🔒
Text Backbone      🔒

Image Projection   ← 학습
Text Projection    ← 학습

이런식으로 projection layer만 먼저 학습할 수도 있다. 그리고 optimizer 또한 학습할 파라미터만 받아줘야한다. 

이부분은 꽤 중요한 부분이기 때문에 다음 블로그에서 집중적으로 다뤄보겠다.

 

전체코드 보러가기([VLM]CLIP_model_implementation)