AI/Deep Learning

[DL] GPU 메모리 사용량 지정 방법

운호(Noah) 2022. 6. 7. 12:50

들어가기 앞서

  • GPU 의 메모리가 부족할 경우,
  • 아래 코드를 통해 GPU 메모리 사용량을 지정할 수 있습니다.

예제 코드

import tensorflow as tf

gpus = tf.config.experimental.list_physical_devices('GPU')

# GPU 메모리를 1024MB (1GB)로 지정합니다.
if gpus: 
      tf.config.experimental.set_virtual_device_configuration(gpus[0],
                      [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])