[실습] 제2장 하드웨어 가속
PyTorch는 CUDA(NVIDIA), MPS(Apple Silicon), ROCm(AMD)을 통해 GPU 가속을 지원한다. 핵심 사용 방식은 장치 변수를 만든 다음, 학습을 시작하기 전에 모델과 데이터를 해당 장치로 이동하는 것이다.
2.1 장치 관리
# 기본 장치 선택
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Apple Silicon (MPS)
device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
# Multi-GPU: specify by index
device = torch.device('cuda:1') # second GPU
# 텐서 및 모델 옮기기
x = x.to(device) # moves tensor
model = model.to(device) # moves all parameters + buffers
x = x.cuda() # shorthand for .to('cuda')
x = x.cpu() # back to CPU
중요:
model.to(device)는nn.Module에 대해서는 제자리 연산이지만 텐서에 대해서는 그렇지 않다. 텐서는x = x.to(device)와 같이 반환값을 다시 할당해야 한다. 참고로 제자리 연산(in-place operate)란 새 객체를 만들지 않고 기존 텐서의 값을 메모리에서 직접 변경하는 연산이다.
2.2 CUDA 유틸리티
| 함수 | 설명 |
|---|---|
torch.cuda.is_available() |
CUDA GPU를 사용할 수 있으면 True를 반환한다. |
torch.cuda.device_count() |
사용할 수 있는 GPU 수를 반환한다. |
torch.cuda.current_device() |
현재 기본 GPU의 인덱스를 반환한다. |
torch.cuda.get_device_name(i) |
i번째 GPU의 이름을 문자열로 반환한다. |
torch.cuda.empty_cache() |
캐시에 남아 있는 미사용 GPU 메모리를 운영체제에 반환한다. |
torch.cuda.memory_allocated() |
현재 GPU에 할당된 메모리의 바이트 수를 반환한다. |
2.3 DataParallel과 DistributedDataParallel
여러 GPU를 사용할 때 PyTorch는 다음 두 가지 다중 GPU 전략을 제공한다.
torch.nn.DataParallel(DP)은 모델을 감싸고 한 컴퓨터의 여러 GPU에 각 배치를 분할한다. 한 줄로 적용할 수 있어 간단하지만 Python GIL 병목이 발생하고 GPU 메모리 사용량이 불균형하다. 본격적인 작업에는 권장하지 않는다.torch.nn.parallel.DistributedDataParallel(DDP)은 GPU마다 하나의 프로세스를 사용한다. 프로세스 간에 기울기를 전체 축소하며 여러 컴퓨터로 확장할 수 있다. GIL 병목이 없지만torch.distributed초기화가 필요하다. 프로덕션 학습에 권장한다.
# DataParallel (simple, not recommended for large scale)
model = torch.nn.DataParallel(model) # wraps the model
# DistributedDataParallel (production standard)
import torch.distributed as dist
dist.init_process_group(backend='nccl')
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])
댓글남기기