[실습] VLA 이전의 시간 — Self-Attention·ViT·CLIP을 직접 굴려보기

pip install torch torchvision transformers Pillow requests

import torch
print(torch.__version__)          # 2.x
print(torch.cuda.is_available())  # GPU 있으면 True

import torch
import torch.nn.functional as F
from math import sqrt

def self_attention(x, W_q, W_k, W_v):
    """
    x: (batch, seq_len, d_model) — 입력 토큰
    W_q, W_k, W_v: (d_model, d_k) — Query·Key·Value 변환 행렬
    반환: (batch, seq_len, d_k) — attention 적용된 새 표현
    """
    Q = x @ W_q  # (batch, seq_len, d_k)
    K = x @ W_k
    V = x @ W_v

    d_k = K.shape[-1]
    scores = Q @ K.transpose(-2, -1) / sqrt(d_k)  # (batch, seq_len, seq_len)
    weights = F.softmax(scores, dim=-1)
    return weights @ V

# 길이 4의 시퀀스, 차원 8로 테스트
torch.manual_seed(0)
x = torch.randn(1, 4, 8)
W_q = torch.randn(8, 8)
W_k = torch.randn(8, 8)
W_v = torch.randn(8, 8)

out = self_attention(x, W_q, W_k, W_v)
print(out.shape)  # torch.Size([1, 4, 8])

from transformers import ViTImageProcessor, ViTForImageClassification
from PIL import Image
import requests

# 사전학습 모델 다운로드 (첫 실행 시 ~350MB)
processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224")

# 테스트 이미지 (COCO val에서)
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)

# 전처리 → 추론
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)

# 결과 해석
predicted_idx = outputs.logits.argmax(-1).item()
print(model.config.id2label[predicted_idx])

probs = F.softmax(outputs.logits, dim=-1)
top5 = probs[0].topk(5)
for p, idx in zip(top5.values, top5.indices):
    print(f"{p.item()*100:5.1f}% — {model.config.id2label[idx.item()]}")

from transformers import CLIPProcessor, CLIPModel

processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")

# 같은 이미지, 임의의 후보 텍스트
candidates = [
    "a photo of a cat",
    "a photo of a dog",
    "a photo of a robot arm",
    "a photo of a sandwich",
]

inputs = processor(
    text=candidates,
    images=image,
    return_tensors="pt",
    padding=True,
)
with torch.no_grad():
    outputs = model(**inputs)

# logits_per_image: (1, len(candidates))
probs = outputs.logits_per_image.softmax(dim=-1)
for label, p in zip(candidates, probs[0]):
    print(f"{p.item()*100:5.1f}% — {label}")

85.2% — a photo of a cat
12.3% — a photo of a dog
 1.8% — a photo of a robot arm
 0.7% — a photo of a sandwich

candidates = [
    "a robot arm picking up an object",
    "a robot arm placing an object",
    "a robot arm idle",
    "an empty workspace",
]