연산 장치 선택 (CPU, CUDA 등)
- PyTorch로 GPU나 NPU 등을 사용 하려는 연산
- 실제 CUDA나 이런 장치를 사용하려면 무조건 이 방식으로 해야한다
장치 확인
if torch.cuda.is_available():
print("CUDA 사용 가능")
else:
print("CUDA 사용 불가능")
실제 연산
# GPU에 경우 cuda()
a = torch.tensor([1.0, 2.0]).cuda()
b = torch.tensor([3.0, 4.0]).cuda()
c = a + b
cuda()
를 호출하면 해당 Tensor가 GPU메모리에 복사가 된다- 사실상
torch
객체에 어디든 이런식으로 사용이 가능하다 - 거의 웬만한 데에서
cuda()
이런식으로 뒤에 호출 해주는 편이 좋다
다른 장치와 혼합 불가
a = torch.tensor([1.0, 2.0]).cuda()
b = torch.tensor([3.0, 4.0]) # CPU Tensor
# 오류 발생: RuntimeError: Expected all tensors to be on the same device
c = a + b
장치를 자동 선택
-
위 방식에 가장 큰 문제는 모든 코드를
cuda()
로 도배 해버리면 cpu만 사용 가능한 상황에서 문제가 발생한다 -
따라서
to(str: device)
함수를 통해 변수값을 사용하여 장치변환 함수를 선택하자if torch.cuda.is_available(): device = "cuda" else: device = "cpu" a = torch.tensor([1.0, 2.0]).to(device)