[Hugging Face] Pre-trained Model Fine-tuning (3. Train Loss vs Validation Loss)

0. Train Loss vs Validation Loss

이전 포스트에서 Training Loss와 Validation Loss 를 비교해 모델이 제대로 fitting이 되었는지 overfitting이 되지는 않았는지 등을 판단 할 수 있다고 했는데 여기서 조금 더 자세히 알아보자.

 

1) Train Loss 와 Validation Loss 의 의미

- Train Loss : 학습에 실제로 사용한 데이터에서 모델이 얼마나 틀리는지 나타낸 지표. 일반적으로 학습이 진행 되면서 점점 잘 맞추기 때문에 loss가 낮아지는 경향을 보임

 - Validation Loss : 학습에는 사용하지 않는 validation 데이터에서 모델이 얼마나 틀리는지 나타낸 지표. 모델이 일반적인 패턴을 잘 학습 했다면 loss가 낮아짐 

 

▲중요한 것은 각각의 절대적인 값이 아닌 두 Loss가 Epoch에 따라 어떻게 움직이는지

 

2) underfitting vs fitting vs overfitting

상태	Train Loss	Validation Loss	특징
Underfitting	높음	높음	충분히 학습되지 않은 상태
Fitting	낮아짐	낮아짐	학습이 잘 진행되는 상태
Overfitting 초기	낮아짐	낮아짐	아직까지는 정상적으로 학습
Overfitting	계속 낮아짐	갑자기 높아짐	Train data에 과도하게 맞춰진 상태

먼저, 표로 확인해 보면 간단히 이렇게 정리 할 수 있다.

당연히 학습된 데이터 내에서 성능이 더 좋을 것이기 때문에 Train Loss < Validation Loss 인 것은 분명하지만 여기서 Validation Loss가 갑자기 커진다거나, Train Loss와 차이가 굉장히 심하게 나는 경우 Ovefitting을 의심 해 봐야한다. Train Data에만 너무 맞춰졌기 때문에 오히려 다른 데이터에서는 제대로 예측을 하지 못하는 경우를 Overfitting 이라고 하는데 그림으로 보면 조금 더 쉽다.

위 그림과 같이 검은색 선으로 학습이 되었다면 Fitting 된 상태이지만 Train data인 파란색 점에만 집중된 초록색 선으로 학습이 되었다면 Overfitting 되었다고 판단하는 것이다. 당연히 자연스럽게 Validation Loss가 올라 갈 수 밖에 없다는 것을 이제 이해 할 수 있을 것이다.

 

1. Loss 비교를 위한 코드 

이전까지 만든 코드는 Train과 validation을 각각 따로 진행 했지만 비교를 위해서는 같이 진행 하는 것이 좋다. 별건 없다 그냥 두개를 하나의 for loop 안에 넣어주면 된다.

train_losses = []
val_losses = []
val_accuracies = []

num_epochs = 3

for epoch in range(num_epochs):

    # Training
    model.train()

    train_loss = 0

    for batch in train_loader:

        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["labels"].to(device)

        optimizer.zero_grad()

        outputs = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels
        )

        loss = outputs.loss

        loss.backward()
        optimizer.step()

        train_loss += loss.item()

    train_loss /= len(train_loader)

    # Validation

    model.eval()

    val_loss = 0
    correct = 0
    total = 0

    with torch.no_grad():

        for batch in val_loader:

            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            labels = batch["labels"].to(device)

            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels=labels
            )

            loss = outputs.loss
            logits = outputs.logits

            val_loss += loss.item()

            predictions = logits.argmax(dim=1)

            correct += (predictions == labels).sum().item()
            total += labels.size(0)

    val_loss /= len(val_loader)
    val_accuracy = correct / total



    train_losses.append(train_loss)
    val_losses.append(val_loss)
    val_accuracies.append(val_accuracy)

    print(
        f"Epoch {epoch+1}"
        f" | Train Loss: {train_loss:.4f}"
        f" | Val Loss: {val_loss:.4f}"
        f" | Val Accuracy: {val_accuracy:.4f}"
    )

이렇게 해주면 결과 값이 Loss를 비교하기에 굉장히 깔끔하게 나오게 된다.

Epoch 1 | Train Loss: 0.4337 | Val Loss: 0.3409 | Val Accuracy: 0.8460
Epoch 2 | Train Loss: 0.2290 | Val Loss: 0.3251 | Val Accuracy: 0.8760
Epoch 3 | Train Loss: 0.1094 | Val Loss: 0.4032 | Val Accuracy: 0.8590

결과 값이 이렇게 나오는데 Epoch3 부터 Val Loss가 다시 높아지는 것을 볼 수 있다. 물론 오래 반복한 것이 아니여서 정확히 Overfitting이다 라고 단정지을 수는 없지만 지금 가장 좋은 모델은 Epoch 2의 모델이기 때문에 Epoch 2 모델을 가지고 Test를 진행 하면 된다.

 

2. 모델 Test

1) best model 저장

이제 우리가 한번도 사용하지 않은 Test 데이터를 사용해 모델을 평가해야한다. 약간 비유를 해 보자면 Training은 혼자 하는 공부, Validation은 모의고사, Test는 실제 시험 쯤으로 생각하면 좋지 않을라나...

아무튼 우리는 Epoch 2 모델을 사용해야 하지만 이미 Epoch 3 까지 진행해버린 상태이기 때문에 다시 Training loop를 돌리면서 최적화된 모델을 저장하는 코드를 넣어줘야한다.

    if val_loss < best_val_loss:
        best_val_loss = val_loss
        torch.save(model.state_dict(), "best_model.pt")
        print("Best model saved!")

이렇게 Val_loss가 가장 작은 모델을 찾아 저장 시키는 코드를 Training Loop 맨 아래에 넣어주면 된다. 

 

2) best model 불러오기

model.load_state_dict(
    torch.load("best_model.pt", weights_only=True)
)

model.to(device)

저장해둔 모델을 이렇게 불러 오면 이제 모델은 Epoch 2 상태가 됐다.

 

3)Test

별거 없다 그냥 Validation 단계를 진행 할 때랑 똑같이 하면 된다

test_loader = DataLoader(
    test_dataset,
    batch_size=16,
    shuffle=False
)

model.eval()

test_loss = 0
correct = 0
total = 0

with torch.no_grad():

    for batch in test_loader:

        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["labels"].to(device)

        outputs = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels
        )

        loss = outputs.loss
        logits = outputs.logits

        test_loss += loss.item()

        predictions = logits.argmax(dim=1)

        correct += (predictions == labels).sum().item()
        total += labels.size(0)

test_loss /= len(test_loader)
test_accuracy = correct / total

print(f"Test Loss: {test_loss:.4f}")
print(f"Test Accuracy: {test_accuracy:.4f}")

Test Loss: 0.5747
Test Accuracy: 0.8470

test 결과 이 모델은 약 84% 확률로 올바르게 분류 했다는 것을 알 수 있다. 

여기서 정확도를 더 높이기 위해서는 Dataset의 크기를 늘리거나, learning-rate를 바꿔보는 등 여러 하이퍼파라미터 들을 바꿔보면서 진행하면 된다.

 

코드보러가기