[Pytorch / Tensorflow] GPU 확인, 특정 GPU 사용 방법

사용가능한 GPU 살펴보는 법 (ubuntu terminal / anaconda terminal 등)

$ nvidia-smi

 

"CUDA_VISIBLE_DEVICES"를 통해 cuda가 볼 수 있는 GPU 제한하기

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3"

별도의 설정을 하지 않으면 cuda는 GPU 0번을 사용하려 한다. GPU:n번 을 사용하려면 번호를 n으로 지정해주면 된다.

위의 예시에서는 GPU:3 이 사용된다.

CPU 강제사용을 원하는 경우 -1로 번호를 선택해주면 된다.

 

Pytorch

import torch
print(torch.cuda.is_available())
# False = cpu
# True = gpu 사용가능

 

사용 예시 :

device = "cuda:3" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")

net = VGG16().to(device)

 

Tensorflow

 

사용가능한 GPU 확인

from tensorflow.python.client import device_lib
device_lib.list_local_devices()