PyTorch에서 사용할 CUDA 버전 지정하기
PyTorch에서 사용할 CUDA 버전 지정하기
다음은 PyTorch에서 사용할 CUDA 버전을 지정하는 방법입니다.
환경 변수 설정
CUDA_VISIBLE_DEVICES 환경 변수를 사용하여 PyTorch가 사용할 GPU 장치를 지정할 수 있습니다. 또한 CUDA_VERSION 환경 변수를 사용하여 특정 CUDA 버전을 선택할 수 있습니다.
# Linux/MacOS
export CUDA_VISIBLE_DEVICES=0
export CUDA_VERSION=11.3
# Windows
set CUDA_VISIBLE_DEVICES=0
set CUDA_VERSION=11.3
torch.cuda.set_device() 함수 사용
torch.cuda.set_device()
함수를 사용하여 특정 GPU 장치를 선택할 수 있습니다. 또한 이 함수는 선택된 장치에 대한 CUDA 버전 정보를 제공합니다.
import torch
# GPU 0 선택
torch.cuda.set_device(0)
# 사용되는 CUDA 버전 확인
print(torch.cuda.get_device_properties().major)
print(torch.cuda.get_device_properties().minor)
torch.cuda.is_available() 함수 사용
torch.cuda.is_available()
함수를 사용하여 시스템에 CUDA가 설치되어 있는지 확인할 수 있습니다.
if torch.cuda.is_available():
print("CUDA is available")
else:
print("CUDA is not available")
torch.cuda.get_device_count()
함수를 사용하여 시스템에 있는 GPU 장치의 개수를 확인할 수 있습니다.
num_gpus = torch.cuda.get_device_count()
print(f"Number of GPUs: {num_gpus}")
참고:
- PyTorch 설치 시 CUDA 버전을 함께 설치해야 합니다.
- 특정 CUDA 버전을 사용해야 하는 경우, Anaconda 또는 Miniconda와 같은 환경 관리 도구를 사용하여 별도의 환경을 만드는 것이 좋습니다.
예제 코드
import torch
# 환경 변수 설정
# Linux/MacOS
export CUDA_VISIBLE_DEVICES=0
export CUDA_VERSION=11.3
# Windows
set CUDA_VISIBLE_DEVICES=0
set CUDA_VERSION=11.3
# GPU 0 선택
torch.cuda.set_device(0)
# 사용되는 CUDA 버전 확인
print(torch.cuda.get_device_properties().major)
print(torch.cuda.get_device_properties().minor)
# CUDA 사용 가능 여부 확인
if torch.cuda.is_available():
print("CUDA is available")
else:
print("CUDA is not available")
# 시스템에 있는 GPU 장치 개수 확인
num_gpus = torch.cuda.get_device_count()
print(f"Number of GPUs: {num_gpus}")
PyTorch에서 사용할 CUDA 버전 지정 대체 방법
torch.cuda.manual_seed_all() 함수 사용
torch.cuda.manual_seed_all()
함수를 사용하여 모든 GPU 장치에 대한 난수 시드를 설정할 수 있습니다. 이 함수는 특정 CUDA 버전을 선택하는 것은 아니지만 특정 버전과 관련된 난수 생성 문제를 해결하는 데 도움이 될 수 있습니다.
torch.cuda.manual_seed_all(seed)
torch.nn.DataParallel 클래스 사용
torch.nn.DataParallel
클래스를 사용하여 여러 GPU 장치에서 모델을 학습할 수 있습니다. 이 클래스는 사용 가능한 모든 GPU 장치를 사용하도록 기본 설정되어 있지만 특정 장치만 사용하도록 설정할 수도 있습니다.
model = torch.nn.DataParallel(model, device_ids=[0, 1])
CUDA_LAUNCH_BLOCKING 환경 변수 설정
CUDA_LAUNCH_BLOCKING
환경 변수를 사용하여 CUDA 런처 동작을 설정할 수 있습니다. 이 환경 변수를 1로 설정하면 CUDA 런처가 모든 작업이 완료될 때까지 차단됩니다. 이는 디버깅에 도움이 될 수 있지만 성능 저하를 초래할 수 있습니다.
# Linux/MacOS
export CUDA_LAUNCH_BLOCKING=1
# Windows
set CUDA_LAUNCH_BLOCKING=1
- 위에서 설명한 대체 방법은 특정 상황에서 유용할 수 있지만 일반적으로는 환경 변수 설정 또는
torch.cuda.set_device()
함수 사용하는 것이 더 좋습니다.
추가 정보
pytorch