pytorch GPU check

일반적인 경우 아래 세 개 command로 pytorch에서 GPU를 사용가능한지 확인이 가능하다.

import torch

torch.cuda.is_available()

torch.cuda.get_device_name(0)

torch.cuda.device_count()

 

각각의 command를 설명하면 다음과 같다.

# GPU를 사용할 수 있는 지 알 수 있다.

torch.cuda.is_available()

→ True

 

# GPU 정보를 반환한다. 아래는 0번 GPU 정보를 반환한다.

torch.cuda.get_device_name(0)

→ 'GeForce GTX 1080 Ti'

 

# 사용가능한 GPU 개수를 반환한다. 아래 예제는 2개의 GPU가 있다는 것을 알 수 있다.

torch.cuda.device_count()

→ 2

# 만약 os 세팅을 통해서 GPU 개수를 줄인다면

# 터미널에서 아래와 같이 환경변수 CUDA_VISIBLE_DEVICES를 0으로 설정한 후 python 실행

CUDA_VISIBLE_DEVICES=0 python                 # python 실행

 

# 환경 변수에 의해서 사용가능한 GPU가 0번만 지정되어 1개로 나타난다.

torch.cuda.device_count()

→ 1

 

# 아래와 같이 사용하는 것은 tourch.cuda.device_count() 결과에 영향을 주지 않는다.

CUDA_VISIBLE_DEVICES=0 python                 # python 실행

 

import os

os.environ['CUDA_VISIBLE_DEVICES]  = "0"

torch.cuda.device_count()

→ 2

 

이외 아래와 같은 명령어들이 있다.

# 항상 0을 return

torch.cuda.current_device()

→ 0

 

# device 정보 반환

torch.cuda.device(0)