PyTorch에서 CUDA 활용법

PyTorch에서 CUDA 활용법

PyTorch는 GPU 가속을 활용하여 대규모 연산을 보다 효율적으로 수행할 수 있도록 지원합니다. 특히 NVIDIA의 CUDA(Compute Unified Device Architecture)를 활용하면 신경망 학습 및 추론 속도를 크게 향상시킬 수 있습니다. 이번 포스팅에서는 PyTorch에서 CUDA를 활용하는 방법을 상세히 살펴보겠습니다.

1. CUDA 개요

CUDA는 NVIDIA에서 개발한 병렬 컴퓨팅 플랫폼이자 API로, GPU에서 일반적인 계산을 수행할 수 있도록 지원합니다. PyTorch는 torch.cuda 모듈을 통해 CUDA를 지원하며, 이를 활용하면 신경망 학습 및 추론을 GPU에서 실행할 수 있습니다.

2. CUDA 사용 가능 여부 확인

PyTorch에서 CUDA가 활성화되어 있는지 확인하려면 아래 코드를 실행하면 됩니다.

import torch

# CUDA 사용 가능 여부 확인
cuda_available = torch.cuda.is_available()
print(f"CUDA 사용 가능 여부: {cuda_available}")

# 사용 가능한 GPU 개수 확인
gpu_count = torch.cuda.device_count()
print(f"사용 가능한 GPU 개수: {gpu_count}")

# 현재 사용 중인 GPU 확인
if cuda_available:
    print(f"현재 사용 중인 GPU: {torch.cuda.get_device_name(0)}")

만약 torch.cuda.is_available()이 True를 반환하지 않는다면 CUDA를 사용할 수 없는 환경입니다. 이 경우, 드라이버 및 PyTorch 설치 상태를 점검해야 합니다.

3. Tensor를 GPU로 이동하기

PyTorch에서 기본적으로 생성된 Tensor는 CPU에서 연산됩니다. 이를 GPU에서 연산하려면 .to() 또는 .cuda() 메서드를 사용해야 합니다.

# CPU에서 생성된 텐서
tensor_cpu = torch.randn(3, 3)
print("CPU Tensor:", tensor_cpu)

# GPU로 이동
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tensor_gpu = tensor_cpu.to(device)
print("GPU Tensor:", tensor_gpu)

이제 tensor_gpu는 GPU에서 연산되며, 이를 다시 CPU로 가져오려면 .cpu()를 사용하면 됩니다.

tensor_back_to_cpu = tensor_gpu.cpu()
print("Back to CPU:", tensor_back_to_cpu)

4. 모델을 GPU로 이동하기

모델을 GPU에서 실행하려면 .to(device)를 사용합니다.

import torch.nn as nn

# 간단한 신경망 모델 정의
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 1)

    def forward(self, x):
        return self.fc(x)

# 모델 생성 및 GPU로 이동
model = SimpleModel().to(device)
print("모델이 위치한 장치:", next(model.parameters()).device)

이렇게 하면 모델의 모든 파라미터가 GPU에 적재됩니다. 이후 입력 데이터도 GPU로 이동해야 모델이 올바르게 동작합니다.

# 입력 데이터도 GPU로 이동
input_data = torch.randn(5, 10).to(device)
output = model(input_data)
print("출력 데이터의 장치:", output.device)

5. GPU 연산 속도 비교

GPU가 CPU보다 얼마나 빠른지 확인하려면 연산 시간을 비교해볼 수 있습니다.

import time

# CPU 연산 시간 측정
cpu_tensor = torch.randn(10000, 10000)
start_time = time.time()
result_cpu = cpu_tensor @ cpu_tensor
end_time = time.time()
print(f"CPU 연산 시간: {end_time - start_time:.5f}초")

# GPU 연산 시간 측정
gpu_tensor = cpu_tensor.to(device)
start_time = time.time()
result_gpu = gpu_tensor @ gpu_tensor
end_time = time.time()
print(f"GPU 연산 시간: {end_time - start_time:.5f}초")

이제 GPU를 사용할 때 연산 속도가 얼마나 향상되는지 직접 확인할 수 있습니다.

6. 여러 GPU 활용 (Data Parallel)

여러 개의 GPU가 있는 경우 DataParallel을 활용하면 모델을 쉽게 병렬 처리할 수 있습니다.

from torch.nn.parallel import DataParallel

# 모델을 여러 GPU에서 병렬 실행
model = SimpleModel()
if torch.cuda.device_count() > 1:
    model = DataParallel(model)

model.to(device)
print("모델이 여러 GPU에서 실행됩니다.")

이렇게 하면 PyTorch가 자동으로 입력 데이터를 여러 GPU에 분배하여 병렬 처리합니다.

7. CUDA 연산 비동기 처리

PyTorch에서 CUDA 연산은 기본적으로 비동기적으로 실행됩니다. 따라서 torch.cuda.synchronize()를 사용하여 연산을 동기화할 수 있습니다.

# 비동기 연산 동기화
start_time = time.time()
result_gpu = gpu_tensor @ gpu_tensor
torch.cuda.synchronize()  # GPU 연산 완료를 기다림
end_time = time.time()
print(f"GPU 연산 (동기화) 시간: {end_time - start_time:.5f}초")

8. 결론

이번 포스팅에서는 PyTorch에서 CUDA를 활용하는 방법을 자세히 살펴보았습니다. CUDA 사용 가능 여부 확인부터 Tensor 및 모델을 GPU로 이동하는 방법, 여러 GPU 활용, 그리고 GPU 연산 속도 비교까지 다루었습니다. GPU를 활용하면 신경망 학습 속도를 크게 향상시킬 수 있으므로, PyTorch를 사용할 때 적극적으로 CUDA를 활용하는 것이 좋습니다.