[Python]GPU를 이용해서 python code 실행하기

유한성·2024년 2월 29일

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

CUDA_VISIBLE_DEVICES는 해당 id의 GPU(들)만을 볼 수 있게 한다는 것으로, 넘겨 받은 GPU(들)을 0번부터로 인식한다.

코드 내부에서 중복해서 선언하는 경우 원하는 index를 제대로 얻지 못할 수 있으므로 중복되지 않도록 주의해야한다

CUDA_DEVICE_ORDER="PCI_BUS_ID" CUDA_VISIBLE_DEVICES="GPU숫자" python 파일명

입력시 해당 파일을 원하는 GPU를 이용하여 실행시킴

watch -d -n 0.5 nvidia-smi

0.5초마다 gpu사용량을 확인하게 해주는 코드이다

0개의 댓글