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)])