파이썬으로 ViT 찍먹하기

CS/Python

파이썬으로 ViT 찍먹하기

munsik22 2026. 8. 18. 13:22

[여기서 이어집니다]

파이썬으로 전이학습 찍먹하기

🧩 전이학습이란?전이 학습(Transfer Learning)은 대규모 데이터셋에서 이미 학습된 모델을 새로운 문제에 재사용하는 방법이다. 일반적으로 CNN(합성곱 신경망, Convolutional Neural Network)을 처음부터 학

munsik22.tistory.com

🧩 ViT란?

ViT(Vision Transformer)는 자연어 처리(NLP)에서 사용되는 트랜스포머 구조를 이미지 인식과 같은 컴퓨터 비전(Computer Vision) 분야에 그대로 적용한 딥러닝 모델 아키텍처다.

CNN과 ViT의 차이

CNN은 작은 필터를 반복적으로 적용해 이미지의 특징을 추출한다.

ViT는 이미지를 여러 개의 작은 패치로 나눈 뒤, 각 패치를 문장의 단어처럼 Transformer에 입력으로 넣는다.

224×224 크기의 이미지를 16×16 크기의 패치로 나눈다면 한 장의 이미지는 (224 ÷ 16)² = 196개의 패치로 분할되고, 각 패치는 벡터로 변환된다.
Transformer의 Self-Attention은 각 패치가 다른 패치와 얼마나 관련이 있는지 계산한다. (귀와 얼굴 형태의 관계, 눈과 코의 위치 관계 등)

💡 CNN이 가까운 영역부터 점진적으로 처리한다면, ViT는 이미지의 서로 떨어진 영역 사이 관계도 직접 학습할 수 있다.

ViT는 이미지 전체를 대표하는 특별한 토큰인 CLS 토큰을 사용한다. Transformer를 통과한 후 CLS 토큰에 이미지 전체 정보가 모이며, 마지막 분류층은 이 값을 사용해 이미지의 클래스를 분류한다.

🧩 HuggingFace Vision Transformer 실습

ViT는 이미지의 먼 영역 사이 관계를 으로 직접 학습할 수 있다는 장점이 있지만, ResNet-18보다 모델이 크기 때문에 더 많은 GPU 메모리와 학습 시간이 필요하다.

모델 불러오기

ViT는 모델이 크고 충분이 학습하려면 많은 데이터가 필요하다. 현재 고양이/강아지 데이터만으로 처음부터 학습하면 과적합될 가능성이 크므로, ImageNet-21k로 사전 학습된 ViT를 사용한다.

from transformers import AutoImageProcessor
from transformers import AutoModelForImageClassification

VIT_CHECKPOINT = "google/vit-base-patch16-224-in21k"
# base: ViT-Base 모델
# patch16: 16×16 패치 사용
# 224: 224×224 이미지 입력
# in21k: ImageNet-21k로 사전 학습

processor = AutoImageProcessor.from_pretrained(VIT_CHECKPOINT)

print("평균:", processor.image_mean)
print("표준편차:", processor.image_std)
print("입력 크기:", processor.size)
# 평균: (0.5, 0.5, 0.5)
# 표준편차: (0.5, 0.5, 0.5)
# 입력 크기: SizeDict(height=224, width=224, longest_edge=None, shortest_edge=None, max_height=None, max_width=None)

만약 ImportError: cannot import name 'ScalingType' from 'torch.nn.functional' (/opt/conda/lib/python3.11/site-packages/torch/nn/functional.py) 에러가 발생한다면, torchao와 torch의 버전이 맞지 않기 때문이다. 현재는 양자화 기능이 필요하지 않기 때문에 torchao를 삭제해주면 된다.

pip uninstall -y torchao

ViT 데이터셋 생성

이전에 CNN에서 만든 것과 동일한 분할 인덱스를 사용해야 공정하게 비교할 수 있다.

vit_train_full = OxfordIIITPet(
    root=DATA_ROOT,
    split="trainval",
    target_types="binary-category",
    transform=vit_train_transform,
    download=False
)

vit_val_full = OxfordIIITPet(
    root=DATA_ROOT,
    split="trainval",
    target_types="binary-category",
    transform=vit_eval_transform,
    download=False
)

vit_test_dataset = OxfordIIITPet(
    root=DATA_ROOT,
    split="test",
    target_types="binary-category",
    transform=vit_eval_transform,
    download=False
)

vit_train_dataset = Subset(vit_train_full, train_indices)
vit_val_dataset = Subset(vit_val_full, val_indices)

ViT는 메모리를 많이 사용하므로 배치 크기를 줄인다.

VIT_BATCH_SIZE = 8

vit_train_loader = DataLoader(
    vit_train_dataset,
    batch_size=VIT_BATCH_SIZE,
    shuffle=True,
    num_workers=NUM_WORKERS,
    pin_memory=torch.cuda.is_available()
)

vit_val_loader = DataLoader(
    vit_val_dataset,
    batch_size=VIT_BATCH_SIZE,
    shuffle=False,
    num_workers=NUM_WORKERS,
    pin_memory=torch.cuda.is_available()
)

vit_test_loader = DataLoader(
    vit_test_dataset,
    batch_size=VIT_BATCH_SIZE,
    shuffle=False,
    num_workers=NUM_WORKERS,
    pin_memory=torch.cuda.is_available()
)

ViT 분류 모델 생성

기존 분류층의 클래스 수와 현재 클래스 수가 다르므로 ignore_mismatched_sizes=True를 사용한다.

id2label = { 0: "cat", 1: "dog" }
label2id = { "cat": 0, "dog": 1 }

vit_model = AutoModelForImageClassification.from_pretrained(
    VIT_CHECKPOINT,
    num_labels=2,
    id2label=id2label,
    label2id=label2id,
    ignore_mismatched_sizes=True
)

vit_model = vit_model.to(device)

분류층은 현재 문제에 맞게 새로 초기화된다. 관련 경고가 출력될 수 있지만 정상이다.

ViT Linear Probing

Linear Probing에서는 Transformer 부분을 동결하고 새 분류층만 학습한다. 사전 학습된 ViT 특징이 고양이와 강아지를 얼마나 잘 구분하는지 확인한다.

# ViT 전체 동결
for parameter in vit_model.parameters():
    parameter.requires_grad = False

# 분류층만 학습 가능하게 설정
for parameter in vit_model.classifier.parameters():
    parameter.requires_grad = True

count_parameters(vit_model)

vit_optimizer = torch.optim.AdamW(
    filter(
        lambda parameter: parameter.requires_grad,
        vit_model.parameters()
    ),
    lr=1e-3,
    weight_decay=1e-4
)

VIT_LINEAR_EPOCHS = 3

vit_linear_history = fit(
    model=vit_model,
    train_loader=vit_train_loader,
    val_loader=vit_val_loader,
    criterion=criterion,
    optimizer=vit_optimizer,
    epochs=VIT_LINEAR_EPOCHS,
    forward_fn=vit_forward
)

plot_history(vit_linear_history, "ViT Linear Probing")

vit_linear_test_loss, vit_linear_test_acc = run_epoch(
    model=vit_model,
    data_loader=vit_test_loader,
    criterion=criterion,
    optimizer=None,
    forward_fn=vit_forward
)

print(
    "ViT Linear Probing Test Accuracy:",
    f"{vit_linear_test_acc * 100:.2f}%"
)

ViT Linear Probing Test Accuracy: 99.84%

ViT 부분 Fine-Tuning

ViT 전체를 Fine-tuning하면 메모리 사용량과 과적합 위험이 커질 수 있다. 따라서 마지막 Transformer 블록 2개만 학습하는 부분 Fine-tuning을 사용한다. 초기 블록은 일반적인 이미지 특징을 유지하고, 마지막 블록은 고양이와 강아지의 귀, 얼굴, 털 등의 특징에 맞게 조정된다.

# 모델의 모든 파라미터 동결
for parameter in vit_model.parameters():
    parameter.requires_grad = False

# 마지막 Transformer 블록 2개 활성화
for parameter in vit_model.vit.encoder.layer[-2:].parameters():
    parameter.requires_grad = True

# 마지막 Layer Normalization과 분류층 활성화
for parameter in vit_model.vit.layernorm.parameters():
    parameter.requires_grad = True

for parameter in vit_model.classifier.parameters():
    parameter.requires_grad = True

count_parameters(vit_model)

vit_backbone_parameters = []
vit_classifier_parameters = []

for name, parameter in vit_model.named_parameters():
    if not parameter.requires_grad:
        continue

    if name.startswith("classifier"):
        vit_classifier_parameters.append(parameter)
    else:
        vit_backbone_parameters.append(parameter)

vit_finetune_optimizer = torch.optim.AdamW(
    [
        {
            "params": vit_backbone_parameters,
            "lr": 1e-5
        },
        {
            "params": vit_classifier_parameters,
            "lr": 1e-4
        }
    ],
    weight_decay=1e-4
)

VIT_FINETUNE_EPOCHS = 3

vit_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    vit_finetune_optimizer,
    T_max=VIT_FINETUNE_EPOCHS
)

vit_finetune_history = fit(
    model=vit_model,
    train_loader=vit_train_loader,
    val_loader=vit_val_loader,
    criterion=criterion,
    optimizer=vit_finetune_optimizer,
    epochs=VIT_FINETUNE_EPOCHS,
    scheduler=vit_scheduler,
    forward_fn=vit_forward
)

plot_history(
    vit_finetune_history,
    "ViT Partial Fine-tuning"
)

vit_finetune_test_loss, vit_finetune_test_acc = run_epoch(
    model=vit_model,
    data_loader=vit_test_loader,
    criterion=criterion,
    optimizer=None,
    forward_fn=vit_forward
)

print(
    "ViT Fine-tuning Test Accuracy:",
    f"{vit_finetune_test_acc * 100:.2f}%"
)

GPU 메모리 부족 오류가 발생하면 VIT_BATCH_SIZE를 줄이거나 마지막 블록 하나만 학습한다.

중간의 vit_finetune_optimizer를 보면 Backbone과 분류층에 서로 다른 학습률을 사용하는 것을 볼 수 있다.

분류층은 새로 만들어졌으므로 상대적으로 큰 학습률(1e-4)을 사용
Transformer는 이미 유용한 특징을 학습한 상태이므로 작은 학습률(1e-5)을 사용

Backbone에 너무 큰 학습률을 사용하면 사전 학습된 유용한 가중치가 빠르게 손상될 수 있는데, 이를 Catastrophic Forgetting이라고 한다.

최종 결과 비교

final_results = {
    "ResNet Scratch": scratch_test_acc * 100,
    "ResNet Linear": linear_test_acc * 100,
    "ResNet Fine-tuning": finetune_test_acc * 100,
    "ViT Linear": vit_linear_test_acc * 100,
    "ViT Fine-tuning": vit_finetune_test_acc * 100
}

print("=" * 45)
print("최종 테스트 결과")
print("=" * 45)

for model_name, accuracy in final_results.items():
    print(f"{model_name:22s}: {accuracy:.2f}%")

plt.figure(figsize=(11, 5))

bars = plt.bar(
    final_results.keys(),
    final_results.values(),
    color=[
        "gray",
        "royalblue",
        "orange",
        "mediumseagreen",
        "tomato"
    ]
)

plt.ylabel("Test Accuracy (%)")
plt.title("Cat vs Dog Classification Results")
plt.ylim(0, 100)
plt.xticks(rotation=20)

for bar, accuracy in zip(bars, final_results.values()):
    plt.text(
        bar.get_x() + bar.get_width() / 2,
        bar.get_height() + 1,
        f"{accuracy:.1f}%",
        ha="center"
    )

plt.tight_layout()
plt.show()

ResNet-18 vs ViT

항목	ResNet-18	ViT-Base
기본 처리 단위	지역적 합성곱	이미지 패치
이미지 전체 관계	층이 깊어지며 확대	Self-Attention으로 직접 계산
이미지에 대한 사전 가정	강함	상대적으로 약함
작은 데이터 학습	비교적 유리	사전 학습이 특히 중요
파라미터 수	약 1170만	약 8600만
연산량	상대적으로 적음	상대적으로 큼
GPU 메모리	적게 사용	많이 사용
주요 특징	지역적 패턴 학습	패치 사이 전역 관계 학습

이번 실습의 경우 ResNet보다 ViT가 성능이 조금 더 높게 나왔지만, 항상 ViT의 성능이 높게 나오는 것은 아니고 적은 데이터셋에서는 ResNet의 정확도가 더 크게 나올 수도 있다.

🧩 핵심 정리

전이 학습 방법 정리

방법	학습 대상	장점	단점
Scratch	전체 모델	사전 모델 의존X	많은 데이터, 시간 필요
Linear Probing	분류층만	빠르고 안정적	새 데이터 적응 한계
Fine-tuning	전체 또는 일부	높은 성능 기대	메모리와 학습시간 증가
데이터 증강	입력 데이터 변환	일반화 성능 향상	너무 강하면 정보 손실
LR Scheduler	학습률 조정	안정적인 최적화	추가 하이퍼파라미터 필요
ViT	Transformer 기반 특징 추출	전역 관계 학습 유리	모델이 크고 계산량 많음

💡 실무에서는 보통 Linear Probing으로 모델과 데이터 파이프라인을 먼저 검증한 후, 작은 학습률로 Fine-tuning을 하는 순서가 안전하다고 한다. (GPT 피셜)

ViT 정리

ViT: 이미지를 패치로 나누어 Transformer로 처리하는 모델
Linear Probing: ViT는 동결하고 분류층만 학습
Fine-tuning: ViT의 일부 또는 전체를 작은 학습률로 추가 학습

'CS > Python' 카테고리의 다른 글

파이썬으로 전이학습 찍먹하기 (0)	2026.08.14
파이썬으로 Attention 찍먹하기 (0)	2026.08.12
파이썬으로 토큰화+임베딩 찍먹하기 (0)	2026.08.12
파이썬으로 MLP 찍먹하기 (0)	2026.08.11
파이썬의 PyTorch 라이브러리에 대해 알아보기 (0)	2026.08.10

현재글파이썬으로 ViT 찍먹하기

MUNSIK BLOG

munsik22 님의 블로그 입니다.

스프링부트, Python, NextJS, MST, springboot, webproxy, 나만무, 크래프톤정글, C++, 레드블랙트리, 프로그래머스, csapp, malloc, 내일배움캠프, pintos, 백준, SWEA, AI, Java, Claude,

Today :
Yesterday :

[Munsik] MUNSIK BLOG [Munsik]