Search

TensorFlow 에서 단일 GPU 메모리 사용량을 제한하는 코드 스니펫

출처
수집시간
2021/09/02 05:03
연결완료
1 more property
def tf_gpu_memory_growing_mode(): """어떤 경우에는 프로세스가 가용한 메모리의 일부에만 할당되도록 하거나 프로세스의 요구량만큼 메모리 사용이 가능할 필요가 있습니다. 텐서플로에서는 이를 위해 두 가지 방법을 제공합니다. 첫 번째 방법은 tf.config.experimental.set_memory_growth를 호출하여 메모리 증가를 허용하는 것입니다. 프로그램이 실행되어 더 많은 GPU 메모리가 필요하면, 텐서플로 프로세스에 할당된 GPU 메모리 영역을 확장합니다. 메모리 해제는 메모리 단편화를 악화시키므로 메모리 해제는 하지 않습니다. 특정 GPU의 메모리 증가를 허용하려면 이 함수를 텐서나 연산 앞에 입력하세요. """ gpus = tf.config.experimental.list_physical_devices('GPU') try: if gpus: tf.config.experimental.set_memory_growth(gpus[0], True) else: print("There are not any physical gpu devices.") except RuntimeError as e: print("This function should call before tensorflow operations called.") # 프로그램 시작시에 메모리 증가가 설정되어야만 합니다 print(e) def tf_gpu_memory_limit_mode(memory_limit=None): """어떤 경우에는 프로세스가 가용한 메모리의 일부에만 할당되도록 하거나 프로세스의 요구량만큼 메모리 사용이 가능할 필요가 있습니다. 텐서플로에서는 이를 위해 두 가지 방법을 제공합니다. 두 번째 방법은 tf.config.experimental.set_virtual_device_configuration으로 가상 GPU 장치를 설정하고 GPU에 할당될 전체 메모리를 제한하는 것입니다. """ assert memory_limit is not None gpus = tf.config.experimental.list_physical_devices('GPU') try: if gpus: tf.config.experimental.set_virtual_device_configuration( gpus[0], [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=memory_limit)]) else: print("There are not any physical gpu devices.") except RuntimeError as e: print("This function should call before tensorflow operations called.") # 프로그램 시작시에 가상 장치가 설정되어야만 합니다 print(e) pass
Python
복사