Tensorflow, Keras, Torch에서 각 버전과 GPU를 사용할 수 있는지를 알아내는 코드입니다.
GPU가 장착된 서버 또는 데스크탑에서도 Nvidia 드라이버와 CUDA를 제대로 설치하지 않아서 사용을 못하는 경우가 많습니다.
GPU가 있고 또 사용할 일이 있을 것 같다면 확인을 하고 시작하는 것이 좋습니다.
# 립러링 프레임워크들의 버전과 GPU 확인
# 1. Tensorflow
import tensorflow as tf
print("Tensorflow version: {}".format(tf.__version__))
print("Tensorflow GPUs: {}".format(tf.test.is_gpu_available(
cuda_only=False,
min_cuda_compute_capability=None
)))
from tensorflow.python.client import device_lib
print("Tensorflow device list: {}".format(device_lib.list_local_devices()))
# 2. Keras
from keras import backend as K
import keras
print("Keras version: {}".format(keras.__version__))
print("Keras GPUs: {}".format(K._get_available_gpus()))
# 3. PyTorch
import torch
print("Torch device count: {}".format(torch.cuda.device_count()))
print("Torch device name: {}".format(torch.cuda.get_device_name(0)))
print("Torch CUDA is available: {}".format(torch.cuda.is_available()))결과는 이런 식으로 나옵니다.
Tensorflow version: 2.6.0
Tensorflow GPUs: True
Keras version: 2.6.0
Keras GPUs: ['/device:GPU:0', '/device:GPU:1']
Torch device count: 2
Torch device name: NVIDIA GeForce RTX 3080
Torch CUDA is available: True