[Python] Python에서 특정 GPU 지정하여 사용하기 - CUDA 설정 가이드

딥러닝, 데이터 과학 프로젝트나 복잡한 연산 작업을 할 때, GPU를 효율적으로 사용해야 하는 상황이 자주 발생합니다. 특히 여러 개의 GPU가 있는 환경에서는 자원 관리를 위해 특정 GPU를 지정하여 사용하는 것이 중요한데요. 이 글에서는 Python 환경에서 특정 GPU를 지정하는 방법을 설명합니다. PyTorch에서의 설정 방법, CUDA_VISIBLE_DEVICES 환경 변수를 이용한 방법 등을 다루어보겠습니다.

 

1. CUDA_VISIBLE_DEVICES 환경 변수로 GPU 지정하기

가장 쉽고 직관적인 방법은 CUDA_VISIBLE_DEVICES 환경 변수를 설정하여 GPU를 선택하는 것입니다. 이 환경 변수는 CUDA가 어떤 GPU만 사용하도록 제어할 수 있게 해주며, 파이썬 프로그램을 실행할 때 GPU 번호를 제한할 수 있습니다.

설정 방법

Linux 또는 Mac 터미널에서 아래와 같이 설정합니다:

CUDA_VISIBLE_DEVICES=0 python train_script.py

 

위 명령어는 GPU 0만 사용하도록 설정합니다. 0을 다른 번호로 변경하여 원하는 GPU를 지정할 수 있습니다.

Windows에서는 터미널에서 환경 변수 설정이 약간 다릅니다. 다음과 같이 실행할 수 있습니다:

set CUDA_VISIBLE_DEVICES=0 && python train_script.py

 

이 명령은 train_script.py를 실행할 때 GPU 0만 인식하도록 합니다.

다중 GPU 지정 방법

두 개 이상의 GPU를 지정하고 싶다면, GPU 번호를 쉼표로 구분하여 설정할 수 있습니다. 예를 들어, GPU 0과 GPU 2를 동시에 사용하려면 다음과 같이 설정하세요:

CUDA_VISIBLE_DEVICES=0,2 python train_script.py

 

이 경우 프로그램은 지정된 두 GPU에 접근할 수 있습니다.

2. Python 코드 내에서 GPU 지정하기

Python 코드 내에서도 특정 GPU를 지정할 수 있습니다. 코드 내에서 환경 변수를 설정하고 GPU를 선택하면, 여러 프로젝트를 실행하는 경우 각 프로젝트에서 다른 GPU를 사용하도록 쉽게 관리할 수 있습니다.

코드 예시

import os
import torch

# 특정 GPU를 지정합니다 (예: GPU 2)
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"  # 장치 순서를 PCI 버스 ID로 설정
os.environ["CUDA_VISIBLE_DEVICES"] = "2"  # GPU 2를 사용하도록 지정

# PyTorch에서 지정된 GPU 사용 여부 확인
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

 

위 코드를 사용하면 GPU 2가 유일하게 사용 가능한 GPU가 되어 torch는 이 GPU만을 인식하게 됩니다.

3. PyTorch에서 GPU 선택과 장치 확인

PyTorch는 다양한 GPU 설정 기능을 제공하여 현재 선택된 장치를 쉽게 확인하고 제어할 수 있습니다. 아래의 코드를 통해 GPU 상태를 확인할 수 있습니다.

import torch

print("총 사용 가능한 GPU 수:", torch.cuda.device_count())
print("현재 사용 중인 GPU 번호:", torch.cuda.current_device())
print("현재 GPU 이름:", torch.cuda.get_device_name(torch.cuda.current_device()))

 

위 코드를 통해 다음 정보를 확인할 수 있습니다:

• 총 사용 가능한 GPU 수: torch.cuda.device_count()는 현재 사용 가능한 GPU 개수를 반환합니다.
• 현재 사용 중인 GPU 번호: torch.cuda.current_device()는 현재 할당된 GPU의 인덱스를 반환합니다.
• GPU 이름: torch.cuda.get_device_name()을 사용하여 GPU 이름을 확인할 수 있습니다.

CUDA_VISIBLE_DEVICES 설정 시 주의할 점

CUDA_VISIBLE_DEVICES 환경 변수를 설정하면 PyTorch는 해당 설정에서 첫 번째로 지정한 GPU를 cuda:0으로 인식하게 됩니다. 예를 들어, CUDA_VISIBLE_DEVICES=2로 설정하면, PyTorch에서 cuda:0은 시스템의 실제 GPU 2에 매핑됩니다. 따라서, GPU를 여러 개 사용하는 경우에는 로컬 인덱스와 실제 인덱스를 혼동하지 않도록 주의해야 합니다.

4. 문제 해결 및 추가 팁

1) GPU 설정이 적용되지 않을 때

가끔 설정이 올바르게 적용되지 않는 경우가 있습니다. 이는 파이썬 환경 내에서 GPU 설정이 변경되지 않는 문제가 발생할 때 주로 나타납니다. 이 경우 Python 인터프리터를 완전히 종료한 후, 환경 변수를 재설정하고 프로그램을 다시 실행해 보세요.

2) 멀티 GPU 활용

딥러닝 학습에서는 큰 모델을 학습할 때 여러 GPU를 사용할 수도 있습니다. torch.nn.DataParallel 또는 torch.nn.parallel.DistributedDataParallel을 사용하여 멀티 GPU 학습을 설정할 수 있습니다. 이를 통해 연산을 여러 GPU에 병렬 분산하여 더 빠르게 학습할 수 있습니다.

---

5. 결론

이번 포스팅에서는 Python에서 특정 GPU를 선택하는 방법에 대해 알아보았습니다. CUDA_VISIBLE_DEVICES 환경 변수와 PyTorch의 torch.device 설정을 통해 시스템의 여러 GPU 중 특정 GPU만을 사용하도록 제어할 수 있습니다. 이러한 설정을 통해 GPU 자원을 효율적으로 사용할 수 있으며, 여러 프로젝트에서 개별적으로 GPU를 관리할 수 있습니다.