Python code killed automatically on AWS EC2 G4 instance - amazon-web-services

Question

I have an AWS EC2 instance g4dn.xlarge and I have to run a python code on it that uses Keras and Tensorflow GPU for the image processing.
I have a large number of images divided into different sets that I need to process using the python code. I'm using python 3.6 and Tensorflow 2.4 and the code runs recursively to process each set of images but after running for some time it gets killed automatically without throwing any error.
The instance I'm using has a single GPU of 16GB and almost all of it is getting used up by the code so I thought maybe it's due to OOM. I looked up the configuration of other available instances on AWS but the others are too large for my requirement.
90% of RAM and CPU is free. I'm just utilizing GPU.
It would be great if there's a way to resolve this issue on the current instance?
Thank you!
P.S. I can only use AWS and not GCP or Azure.

Answer 1

I figured it out. Apparently, TensorFlow uses up all the GPU memory available on the system whether it needs it or not which in this case was 16GB. Due to the overutilisation of GPU, the python script was getting killed automatically after some time.
To make your script use only the required memory add the "Limiting GPU memory growth" feature available on the TensorFlow documentation.
https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth
gpus = tf.config.list_physical_devices('GPU')
if gpus:
# Restrict TensorFlow to only use the first GPU
try:
tf.config.set_visible_devices(gpus[0], 'GPU')
logical_gpus = tf.config.list_logical_devices('GPU')
print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPU")
except RuntimeError as e:
# Visible devices must be set before GPUs have been initialized
print(e)
You need to add the above script in your main function or if you don't have a main function add it after you have imported all the libraries.
After adding this to the python code only 50% of the GPU memory is being utilised(which was the actual requirement) and the code is running fine without any error or being killed.