Limit GPU Virtual Memory for Tensorflow | 텐서플로우 GPU 메모리 부족으로 학습중단될 경우 제한방법

인공지능을 학습하다보면 자주 메모리문제를 보게 되는데요. 단순히 사용하는 것만으로 설정할 경우 메모리를 한계치 이상으로 사용하게되거나 메모리할당이 안되거나 하면서 학습이 중지 되는 경우가 있습니다. 따라서 여기서는 메모리의 사용량을 설정하는 방법입니다.

When learning AI, we often encounter memory problems. If you simply set GPU and use it, you may end up using more memory than your laptop has, or you may not be able to allocate memory, and learning stops. Therefore, this section explains how to set the memory limit. 

 

간단하게 아래와 같이 GPU가 실제로 설정되어 있는지 확인하고 설정하는 방법입니다. 

Simply, here's checking the GPUs able to use and set

gpu_devices = tf.config.experimental.list_physical_devices('GPU')
if gpu_devices:
    print('Using GPU')
    tf.config.experimental.set_memory_growth(gpu_devices[0], True)
else:
    print('Using CPU')

 

아래와 같이 VirtualDeviceConfiguration 에 memory_limit 를 설정 하면 됩니다.

You can set the source code to limit GPU usage like below.

gpu_devices = tf.config.experimental.list_physical_devices('GPU')
if gpu_devices:
    print('Using GPU')
    # limit GPU Virtual Memory for 5GB
    tf.config.experimental.set_virtual_device_configuration(gpu_devices[0],
        [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024*6)])
    # tf.config.experimental.set_memory_growth(gpu_devices[0], True) # can't use when set_virtual_device
else:
    print('Using CPU')

반응형

기본적으로 텐서플로우에서는 물리적 메모리제한을 초과하지 않도록 설계되어 있으나 메모리문제는 자주 발생합니다. 위의 설정은 물리적 GPU가 2개 이상있을 경우 유용하지만 1개일 경우에도 메모리 사용제한이나 사용량제어에 도움이 됩니다.(실제로 단순 GPU설정으로 오류가 많이 발생했었을때 위의 가상디바이스 설정으로 해결된 경우가 많음) 하지만 확실한건, 모든 메모리 문제를 해결할 수 있는 건 아닙니다. 
Basically, Tensorflow has been set not to use over the physical memory limit. However, memory problem occurs often during AI calculation. Of course, the above method is useful if you have an additional physical GPU. However, 1 GPU also can be used to solve memory problems by the above setting. Not always tho.

 

배치를 좀 더 분할 한다거나 하는 메모리 사용 테크닉도 필요할 수도 있습니다~

You may need some technics such as reducing the batch size to optimize the memory usage.