[꼬꼬무] - PyTorch 에서 메모리 부족할때 해결하는 방법! (Autocast, GradScaler)

Autocast와 GradScaler

딥러닝 모델을 학습할 때, 특히 GPU를 사용할 경우 메모리 효율성과 연산 속도를 개선하는 것이 매우 중요하다. 이때 PyTorch에서 제공하는 Autocast와 GradScaler는 메모리 사용량을 줄이고 학습 성능을 최적화하는 데 효과적이다. 이 글에서는 Autocast와 GradScaler의 원리를 설명하고, 이를 활용한 코드와 적용 사례를 알아보자.

---

Autocast의 원리

Autocast는 PyTorch에서 제공하는 기능으로, Mixed Precision Training(혼합 정밀도 학습)을 구현한다.

 

그러면 Mixed Precision Training 이게 무엇인가??

딥러닝 연산에서 일부를 FP32(32-bit floating point) 대신 FP16(16-bit floating point) 정밀도로 수행하여 메모리 사용량을 줄이고 연산 속도를 높이는 기술이다. FP16 연산은 더 적은 메모리를 사용하고 GPU의 Tensor Core를 활용하여 연산 성능을 극대화할 수 있다.

---

FP32, FP16은 무엇인가??

FP16과 FP32는 컴퓨터 연산에서 사용하는 부동소수점(Floating Point) 표현 방식의 하나로, 숫자를 저장하고 연산하는 데 필요한 비트의 크기와 관련이 있다. 이 개념은 딥러닝 연산에서 정밀도(Precision와 메모리 사용량에 영향을 미치며, 연산 속도와도 밀접한 관련이 있다.

---

부동소수점(Floating Point)은 무엇인가?

우선, 컴퓨터는 2진법으로 모든 데이터를 나타낸다. 

부동소수점은 컴퓨터가 실수를 표현하는 방법 중 하나이다. 

아래 예시를 통해 이해해보자

고정소수점

고정소수점은 숫자를 표현할 때 소수점의 위치를 고정하여 정수 부분과 소수 부분으로 나누는 방식이다.

• bit를 반으로 나눠서 절반은 정수, 절반은 소수 부분에 할당한다.
• 정수 부분이 큰 숫자, 소수 부분이 정밀한 숫자를 나타낼 수 없다.
• 정수 부분을 늘리면 소수 부분이 줄어들고, 소수 부분을 늘리면 정수 부분이 줄어들게 된다.

예시: 소수점이 "정수 4자리 + 소수 2자리"로 고정된 경우

1234.56은 표현 가능.

하지만 0.0001이나 1234567 같은 값은 표현 불가능(범위 초과).

부동소수점

• 모든 숫자를 1.xxxxxx 형식으로 나타낸다.
• 첫 번째 비트는 양수, 음수를 구분하는데 사용한다.
• 그 다음 8비트로 소수점이 몇 칸 움직일 지를 나타낸다. (127과의 차이)
• 나머지 23자리에 소수점이 움직인 결과에서 소수점 뒤로 오는 부분들을 채워넣는다.

예시: "1.XXXXX × 2 ^ 지수" 형태

0.0001은 1.0 × 10 ^ -4로 표현 가능.

1234567은 1.234567 × 10 ^ 6으로 표현 가능.

숫자의 크기에 따라 소수점 위치를 조정해 다양한 값 표현이 가능하다.

 

따라서

Pytorch 에서는 기본적으로 FP32(32-bit 부동소수점) 를 사용한다.

고정 소수점을 사용하면 정수, 소수 부분으로 나누기 때문에 메모리 낭비가 발생할 가능성이 높다 

(물론 1개 1개 변수에 데이터 타입을 지정하면 효율이 높아질 수 도 있다..)

그래서 Pytorch 에서는 소수점 위치를 유동적으로 조정하는 부동소수점을 사용해서, 메모리를 더 효율적으로 사용할 수 있다.

---

 

 

다시, Mixed Precision Training 이란?

딥러닝 연산에서 일부를 FP32(32-bit floating point) 대신 FP16(16-bit floating point) 정밀도로 수행하여 메모리 사용량을 줄이고 연산 속도를 높이는 기술이다. FP16 연산은 더 적은 메모리를 사용하고 GPU의 Tensor Core를 활용하여 연산 성능을 극대화할 수 있다.

---

 

그러면 Pytorch 에서 메모리를 절약하기 위해 어떻게 FP16 을 활용해서 적용할까??

 

PyTorch는 Autocast와 GradScaler라는 도구를 제공하여 FP16을 간단하고 안정적으로 적용할 수 있다.

 

Autocast: FP16과 FP32의 혼합 사용

• Autocast는 모델의 연산을 FP16과 FP32를 혼합하여 수행한다.
• FP16으로 처리 가능한 연산은 FP16으로 실행하고, 정밀도가 필요한 연산(예: Softmax, Logarithm 등)은 FP32로 자동 변환한다.
• 이를 통해 메모리 사용량을 줄이면서도 모델의 학습 안정성을 유지할 수 있다.

GradScaler: Gradient 문제 해결

FP16은 정밀도가 낮아, Gradient 값이 너무 작아질 경우 Underflow(값이 0으로 처리됨) 문제가 발생할 수 있다. 이를 해결하기 위해 PyTorch는 GradScaler를 제공한다.

• GradScaler는 Gradient 값을 스케일 업(확대)하여 Underflow를 방지한다.
• Optimizer 업데이트 직전에 Gradient를 다시 원래 크기로 조정(스케일 다운)하여 학습 안정성을 보장한다.

import torch
from torch.cuda.amp import autocast, GradScaler

# 모델 및 데이터 로더 초기화
model = MyModel().cuda()  
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
criterion = torch.nn.CrossEntropyLoss()  # 손실 함수
scaler = GradScaler()  # GradScaler 초기화

# 데이터 로더
dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 학습 루프
for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    
    for images, labels in dataloader:
        images, labels = images.cuda(), labels.cuda()
        
        optimizer.zero_grad()  # Optimizer 초기화
        
        # Autocast로 FP16 적용
        with autocast():
            outputs = model(images)  # 모델 출력 (FP16 연산)
            loss = criterion(outputs, labels)  # 손실 계산
        
        # GradScaler로 Gradient 계산 및 업데이트
        scaler.scale(loss).backward()  # 손실 스케일 업 후 역전파
        scaler.step(optimizer)  # Optimizer 업데이트
        scaler.update()  # 스케일러 업데이트
        
        total_loss += loss.item()
    
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss / len(dataloader):.4f}")

 

---

 

Gradient Underflow 를 방지하기 위해 GradScaler 를 사용한다는게 무슨말인가??

 

Gradient Underflow는 Gradient 값이 너무 작아져서 부동소수점(Floating Point)이 표현할 수 있는

최소값보다 작아질 때 발생한다.

쉽게 말하면,

 

• Gradient 값이 0.0000002(아주 작음)이라고 가정하자.
• FP16으로 연산하면 이 값은 표현할 수 없기 때문에 0으로 저장된다.
• 결과적으로, 학습이 제대로 이루어지지 않는다.

 

그래서 GradScaler는 이런 Underflow 문제를 해결하기 위해 Gradient 값을 확대(Scale Up)하여

작은 값도 표현할 수 있도록 도와준다.

아래 예시로 설명을 하면,

1. Gradient를 스케일 업:
   • Gradient 값에 큰 숫자(예: 1000)를 곱해서 값을 크게 만든다.
   • 예: 0.00002×1000 = 0.02. 이제 FP16으로도 표현 가능해진다.
2. Optimizer 업데이트 후 원래 크기로 스케일 다운:
   • 가중치를 업데이트한 뒤, Gradient를 다시 원래 값으로 나눈다.
   • 예: 0.02 /1000=0.00002. 모델의 학습 과정에서 실제 Gradient 값을 복원한다.

 

GradScaler 사용의 장점은 무엇일까?

1. Underflow 방지:
   • Gradient 값이 0으로 처리되는 문제를 해결하여 학습이 멈추는 것을 방지한다.
2. FP16의 메모리 효율을 유지:
   • Gradient 값을 확장하고 조정하면서도 FP16을 계속 사용할 수 있어 메모리를 절약할 수 있다.
3. 모델 학습 안정화:
   • Gradient 값을 조정하여 학습이 보다 안정적으로 진행된다.

 

비유로 정리를 하면 

Gradient Underflow와 GradScaler를 비유로 설명하면 다음과 같다:

Gradient Underflow은

아주 작은 글씨를 쓴다고 생각하자. 글씨가 너무 작아서 돋보기를 써도 보이지 않으면 "0"으로 간주하는 상황이다.

 

GradScaler 로

Gradient 값을 확대하는 것은 마치 글씨를 크게 써서 돋보기를 통해 잘 보이게 만드는 과정이다.

Optimizer 업데이트 후 원래 글씨 크기로 되돌리는 과정이 GradScaler의 역할이다.

 

그래서 

PyTorch에서 FP16을 활용하려면 GradScaler와 Autocast를 반드시 함께 사용하는 것이 효과적이다!

 

끝 이상입니다.

감사합니다.