1 분 소요

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])

댓글남기기