[Pytorch] Pytorch로 GPU 사용 가능한지 확인하기

Pytorch를 통해 GPU를 사용할 수 있는지, 그 수는 몇인지 확인하는 방법

 

import torch

# CUDA 사용한지 여부 확인: True or False
print(torch.cuda.is_available())

# GPU 수 확인
print(torch.cuda.device_count())

# 현재 사용하는 GPU 번호 확인
print(torch.cuda.current_device())

# 0번 GPU 이름 확인: ex) GeForce GTX 950M
print(torch.cuda.get_device_name(0))

 

 

Reference

https://stackoverflow.com/questions/48152674/how-to-check-if-pytorch-is-using-the-gpu

How to check if pytorch is using the GPU?