Pytorch에서 GPU 사용

파이토치에서 GPU를 사용하기 위한 과정을 정리

CUDA 및 Pytorch설치

gpu를 사용하기 위해선 CUDA toolkit과 cuDNN이 필요하다.

GPU에 맞는 CUDA Toolkit 버전 확인

링크1에서 gpu의 Compute Capability를 확인
링크2에서 gpu의 Compute Capability에 맞는 CUDA Toolkit과 cuDNN의 버전 확인
파이토치 최신 버전과 파이토치 이전 버전에서 사용하려는 CUDA Toolkit을 지원하는지 확인

파이토치 최신버전을 사용한다면 여기서 옵션을 선택한 후
명령어를 복사 후 실행

설치

원하는 위치에서
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128
파이토치와, CUDA와 cuDNN이 설치된다.

확인

import torch
print('Torch :', torch.__version__)
print( 'CUDA :' , torch.version.cuda)
print('cuDNN :',torch.backends.cudnn.version())

if torch.cuda.is_available():
    print("GPU 사용 가능!")
    print(f"사용 가능한 GPU 개수: {torch.cuda.device_count()}")
    print(f"첫 번째 GPU 이름: {torch.cuda.get_device_name(0)}")
else:
    print("GPU를 사용할 수 없습니다.")

GPU 사용 가능!
사용 가능한 GPU 개수: 1
첫 번째 GPU 이름: NVIDIA GeForce RTX 4060 Laptop GPU
'2.7.1+cu128'
Torch : 2.7.1+cu128
CUDA : 12.8
cuDNN : 90701
GPU 사용 가능!
사용 가능한 GPU 개수: 1
첫 번째 GPU 이름: NVIDIA GeForce RTX 4060 Laptop GPU

Torch, CUDA, cuDNN의 버전이 제대로 설치되었고 GPU가 사용 가능한 것을 확인할 수 있다.