특정 GPU에 memory 할당하기

Hα ყҽσɳɠ·2021년 8월 1일
0

  • 여러 명이 GPU를 사용할 때
  • 특정 실험을 특정 GPU에 올리고자 할 때

위와 같은 경우에 특정 GPU에 memory를 할당해야 하는데, 다음 명령어를 이용하면 된다.

0번 GPU(TITAN RTX)를 사용하고 싶은 경우,

🛠 1. src code 내에 선언

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

주로 train 코드 내에 선언하거나, main내에 train() 위에 선언하여 사용한다.

🛠 2. 실행 시 선언

CUDA_VISIBLE_DEVICES=0 python main.py

프로그램을 run할 때, 할당하고 싶은 GPU의 번호를 선언해줄 수도 있다.

profile
𝑯𝒐𝒏𝒆𝒔𝒕𝒚 𝑰𝒏𝒕𝒆𝒈𝒓𝒊𝒕𝒚 𝑬𝒙𝒄𝒆𝒍𝒍𝒆𝒏𝒄𝒆

0개의 댓글