torch.autocast는 PyTorch에서 **자동 혼합 정밀도(Automatic Mixed Precision, AMP)**를 매우 쉽게 사용하도록 도와주는 **컨텍스트 매니저(context manager)**입니다.
간단히 말해, torch.autocast는 모델의 훈련 및 추론 속도를 높이고 GPU 메모리 사용량을 줄이기 위해 사용됩니다.
torch.autocast는 PyTorch에서 **자동 혼합 정밀도(Automatic Mixed Precision, AMP)**를 매우 쉽게 사용하도록 도와주는 **컨텍스트 매니저(context manager)**입니다.
간단히 말해, torch.autocast는 모델의 훈련 및 추론 속도를 높이고 GPU 메모리 사용량을 줄이기 위해 사용됩니다.
1. 왜 autocast가 필요한가요?
딥러닝 모델은 기본적으로 32비트 부동 소수점(float32 또는 FP32)으로 연산을 수행합니다. 이는 높은 정밀도를 보장하지만, 다음과 같은 단점이 있습니다.
- 느린 속도: 연산량이 많습니다.
- 높은 메모리 사용량: 각 숫자가 32비트(4바이트)를 차지합니다.
최신 GPU(NVIDIA의 경우 ‘Tensor Cores’가 탑재된)는 float16(FP16)이나 bfloat16(BF16) 같은 16비트 연산을 훨씬 빠르게 처리할 수 있습니다. 16비트 연산은 다음과 같은 이점을 제공합니다.
- 빠른 속도: 동일 시간당 더 많은 연산이 가능합니다. (이론적으로 2배 이상)
- 낮은 메모리 사용량: 필요한 메모리가 절반으로 줄어듭니다. (더 큰 모델, 더 큰 배치 크기 가능)
하지만 모든 연산을 FP16으로 바꾸면, 정밀도가 너무 낮아져서 **수치적 불안정성(numerical instability)**이 발생할 수 있습니다. 특히 그래디언트(gradient)가 0에 너무 가까워져 사라지는 ‘언더플로우(underflow)’ 문제가 생기기 쉽습니다.
2. autocast는 무엇을 “자동”으로 하나요?
autocast는 이 문제를 해결하기 위해 “혼합 정밀도” 전략을 사용합니다. 즉, 연산에 따라 FP32와 FP16을 자동으로 선택합니다.
FP16(또는BF16)으로 실행되는 연산:- 행렬 곱(
matmul), 컨볼루션(conv) 등 속도 향상에 큰 영향을 주는 연산. - GPU의 텐서 코어를 활용하여 속도를 극대화합니다.
- 행렬 곱(
FP32로 유지되는 연산:Softmax, 손실 함수(Loss function) 계산, 배치 정규화(Batch Normalization) 등 수치적 안정성이 중요한 연산.- 정밀도를 유지하여 모델이 안정적으로 학습되도록 합니다.
autocast 컨텍스트(with torch.autocast(...)) 내에서 실행되는 연산들은 PyTorch가 알아서 “이 연산은 FP16으로 해도 괜찮아” 또는 “이 연산은 FP32로 유지해야 해”라고 판단하여 실행합니다.
import torch
#build_sam2: SAM2 모델의 아키텍처를 구축하는 함수이다.
from sam2.build_sam import build_sam2
#SAM2ImagePredictor: SAM2모델을 사용해 이미지에 대한 예측(세그멘테이션)을 알기 쉽게 수행할 수 있도록 도와주는 래퍼(wrapper) 클래스이다.
from sam2.sam2_image_predictor import SAM2ImagePredictor
#checkpoint: 미리 학습된 모델의 가중치(weight) 파일 경로입니다. (.pt 파일)
checkpoint = "../checkpoints/sam2.1_hiera_large.pt"
#model_cfg: 모델의 구조를 정의하는 설정(configuration) 파일 경로입니다. (.yaml 파일)
model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
#build_sam2(model_cfg, checkpoint): 설정 파일(model_cfg)을 기반으로 SAM2 모델의 뼈대를 #만들고, 그 뼈대에 학습된 가중치(checkpoint)를 불러옵니다. 이렇게 완성된 모델 객체가 반환됩니다.
#SAM2ImagePredictor(...): 방금 로드한 SAM2 모델을 SAM2ImagePredictor 클래스에 전달하여, 추론을 위한 predictor 객체를 생성합니다. 이 predictor를 사용해 이미지를 입력하고 마스크를 출력받게 됩니다.
predictor = SAM2ImagePredictor(build_sam2(model_cfg, checkpoint))
#torch.inference_mode():
#"추론 모드"를 활성화합니다.
#PyTorch가 그래디언트(gradient)를 계산하거나 추적하지 않도록 설정합니다.
#학습이 아닌 순수 예측(추론) 시에는 이 모드를 사용하는 것이 메모리를 절약하고 속도를 향상시키는 데 필수적입니다. (torch.no_grad()와 유사하지만 더 빠릅니다.)
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
predictor.set_image(rf"/root/Coding/sam2/notebooks/images/cars.jpg")
masks, _, _ = predictor.predict("")torch.autocast("cuda", dtype=torch.bfloat16):
- **자동 혼합 정밀도(AMP)**를 활성화합니다.
"cuda": 이 연산을 GPU에서 수행하도록 지정합니다.dtype=torch.bfloat16: 연산 시 사용할 16비트 데이터 타입을bfloat16으로 지정합니다.- 효과:
bfloat16은float16보다 수치적 안정성이 높으면서도float32보다 훨씬 빠릅니다. 이with블록 안에서 실행되는 모델 연산(주로 행렬 곱, 컨볼루션 등)은 자동으로bfloat16으로 수행되어 추론 속도가 매우 빨라집니다.
masks, _, _ = predictor.predict(""):
predictor에게 예측을 수행하라고 명령합니다.""(빈 문자열)을 프롬프트로 전달했는데, 이는 SAM 모델에서 “이미지 내의 모든 것을 자동으로 분할하라(segment everything)”는 의미일 가능성이 큽니다.predictor.predict는 보통(마스크, 점수, 로짓)튜플을 반환합니다.masks, _, _: 첫 번째 반환 값인 **masks**만 변수에 저장하고, 나머지 두 값(아마도 품질 점수 등)은_를 사용해 무시(discard)합니다.
요약
이 코드는 sam2.1_hiera_large라는 강력한 SAM2 모델을 로드한 뒤, **torch.inference_mode**와 **torch.autocast (bfloat16)**를 이용해 GPU에서 매우 빠르고 효율적으로 cars.jpg 이미지에 대한 자동 세그멘테이션 마스크를 생성하는 과정입니다.
결과적으로 masks 변수에는 cars.jpg 이미지의 객체들이 분할된 마스크 데이터가 (아마도 텐서 형태로) 담기게 됩니다.
