[파이썬] PyTorch GPU 가속 사용 방법

07 Sep 2023

PyTorch

PyTorch는 GPU를 활용하여 신경망 모델의 학습과 추론 속도를 향상시킬 수 있습니다. GPU를 사용하여 연산을 수행하면 CPU보다 훨씬 빠른 속도로 모델을 실행할 수 있습니다. 이 글에서는 PyTorch에서 GPU를 가속화하는 방법을 알아보겠습니다.

GPU 지원 확인

먼저, 시스템에서 GPU가 제대로 지원되는지 확인해야 합니다. CUDA와 함께 사용되는 NVIDIA GPU를 사용하는 경우가 가장 일반적입니다. 다음 명령으로 GPU 지원을 확인할 수 있습니다.

import torch

if torch.cuda.is_available():
    print("GPU를 사용할 수 있습니다.")
else:
    print("GPU를 사용할 수 없습니다.")

장치 설정

GPU를 사용하기 위해선 Torch Tensor를 GPU로 이동시켜야 합니다. 기본적으로 Tensor는 CPU 상에서 생성되며, 이를 GPU로 이동시켜 가속을 활용할 수 있습니다. 다음 예제를 통해 GPU 장치 설정 방법을 확인해보겠습니다.

import torch

# GPU 사용 가능여부 확인
if torch.cuda.is_available():
    device = torch.device("cuda")  # CUDA 장치 객체 생성
    print("GPU를 사용합니다.")
else:
    device = torch.device("cpu")  # CPU 장치 객체 생성
    print("GPU를 사용할 수 없습니다.")

# Tensor를 GPU로 이동
x = torch.tensor([1., 2.]).to(device)

모델에 GPU 적용

이제 모델에 GPU를 적용하는 방법을 살펴보겠습니다. 모델을 GPU로 보내기 위해서는 모델의 파라미터들을 GPU 메모리로 이동시켜야 합니다. 다음은 모델을 GPU로 이동시키는 예제입니다.

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.linear = nn.Linear(2, 1)

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

model = MyModel()

# GPU 사용 가능여부 확인
if torch.cuda.is_available():
    model.to(device)  # 모델을 GPU로 이동
    print("모델을 GPU로 이동시켰습니다.")
else:
    print("GPU를 사용할 수 없어 모델을 CPU에 유지시킵니다.")

데이터 로드와 GPU

마지막으로 GPU 상에서 모델과 함께 데이터를 로드하는 방법을 알아보겠습니다. 데이터를 GPU로 이동하기 위해서는 데이터 로더에서 Tensor를 GPU로 이동시켜야 합니다. 다음은 데이터 로드와 GPU 이동 예제입니다.

import torch
from torch.utils.data import DataLoader

# 데이터셋 및 데이터로더 생성
dataset = ...
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 데이터 GPU로 이동
for inputs, labels in dataloader:
    inputs = inputs.to(device)
    labels = labels.to(device)

    # 모델 연산 수행
    outputs = model(inputs)

마치며

이번 글에서는 PyTorch에서 GPU를 사용하여 모델의 학습과 추론 속도를 향상시키는 방법에 대해 알아보았습니다. GPU를 사용하면 CPU와 비교하여 훨씬 빠른 속도로 성능을 높일 수 있습니다. 이제 여러분은 PyTorch에서 GPU 가속을 적용하여 더욱 빠른 딥러닝 모델을 개발할 수 있을 것입니다!