[Python / Linux] GPU 메모리 비우기, 프로세스 전부 종료하기

---

GPU 메모리 비우기, 프로세스 전부 종료하기

pytorch나 tensorflow와 같은 인공지능 프레임워크를 사용하다보면, 중간에 예기치 않게 프로그램이 종료될 때가 있는데, 이 때 문제점이 data_loader의 num_worker 수만큼 발생한 서브프로세스들이 부모 프로세스가 종료되었다는 사실을 인지하지 못하고 계속 작업을 수행하는 경우가 발생한다.

 

즉, 프로그램이 종료되었는데도 불구하고, GPU 메모리는 모두 반환되지 않는 상태가 종종 발생한다.

매우 단순한 해결 방법을 알아보자!

 

해결 방안

---

for i in $(lsof /dev/nvidia1 | grep python | awk '{print $2}' | sort -u); do kill -9 $i; done

위 명령어를 통해서, nvidia gpu를 점유하고 있는 모든 python 프로세스를 종료시킬 수 있다.

multi-gpu 환경이라면, nvidia1을 nvidia2로 바꾸는 등 인덱스에 맞게 여러번 실행시키면 된다.

 

일반적으로 torch.dist와 같은 multi-gpu training의 경우엔 nvidia1에 해당하는 프로세스만 꺼줘도 모두 정상 반환되는 것을 확인할 수 있었다.