PyTorch에서 GPU setting

Pytorch 버전 확인

torch.__version__

 

CUDA 버전 확인

torch.version.cuda

 

cudnn 버전 확인

torch.backends.cudnn.version()

 

CUDA 사용 가능 확인

torch.cuda.is_available()

 

GPU 갯수 확인

torch.cuda.device_count()

 

현재 사용 GPU 확인

torch.cuda.current_device()

 

GPU가 여러개인 경우 원하는 GPU 선택하기

 

nvidia-smi

를 통해 GPU 정보를 볼 수 있다.

 

device = torch.device(f'cuda:{GPUtoUse}' if torch.cuda.is_available() else 'cpu')

GPUtoUse 부분에 숫자를 바꾸어 몇번째 GPU를 사용할지 지정가능하다.

 

예시.

0번 GPU 사용 시

device = torch.device(f'cuda:0' if torch.cuda.is_available() else 'cpu')

 

 

 

===============================================================

#### GPU SETTING ####
print("----------------------------------------------------------------")
print(f"PyTorch Version: {torch.__version__}")
print(f"CUDA Version: {torch.version.cuda}")
print(f"cudnn Version: {torch.backends.cudnn.version()}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"Count of GPUs: {torch.cuda.device_count()}")
print("----------------------------------------------------------------")
GPUtoUse = input("select GPU to use: ")

net = Unet.UNet(in_channels=1, n_classes=2, padding=True)
torchsummary.summary(net, (1, 256, 256), device='cpu')
device = torch.device(f'cuda:{GPUtoUse}' if torch.cuda.is_available() else 'cpu')
print(f"Using GPU model: {torch.cuda.get_device_name(device)}, Selected GPU: {GPUtoUse}")
net = net.to(device=device)

time.sleep(1.5)