본문 바로가기

Data-science/deep learning

[tensorflow] 메모리 과부하시 해결법

728x90
import tensorflow
physical_devices = tensorflow.config.list_physical_devices('GPU')

try:
    tensorflow.config.experimental.set_memory_growth(physical_devices[0], True)
    tensorflow.config.experimental.set_memory_growth(physical_devices[1], True)
except:
    # Invalid device or cannot modify virtual devices once initialized.
    pass

GPU가 2개일 경우 위 코드를 쓴다.

하나일 경우 [0]하나만 쓰자!

 

이렇게할 경우 gpu를 한 번에 다 잡고 시작하는게 아니라, 네트워크에 해당하는 메모리만큼만 잡으며 시작한다.