I have a set of images that I need to run an inference on. What I am trying to do is spawn 4 workers and each worker has 1/4 access to GPU's memory. I am getting OOM even if I restrict the memory to 1/8 per worker. Below is some of the code:
def process_image(doc_paths):
model = models.Model(my_model)
for path in doc_paths:
...do some work...
model.do_predict()
return 0 or 1
if __name__ == "__main__":
num_processes = 4
process_pool_executer = concurrent.futures.ProcessPoolExecutor(num_processes)
# Split documents into `num_processes` equal parts, so in our case -> 4 equal parts
for i, paths in enumerate(sub_file_paths):
print(f'==== Processing path {i} ====')
f = process_pool_executer.submit(process_image, paths)
futures.append(f)
for future in concurrent.futures.as_completed(futures):
r = future.result()
results.append(r)
Within the do_predict() function, I import tensorflow and have tried both these memory allocation options:
# OPTION 1
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
tf.config.experimental.set_virtual_device_configuration(gpus[0], [
tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])
except RuntimeError as e:
print(e)
# OPTION 2
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
Both give me OOM errors. According to nvidia-smi my GPU has 8081 MiB free so allocating 1024 MB for each process should be fine. So I must be doing something or thinking of something in the wrong way. Any thoughts?