Supporting multiprocessing inference on single GPU for Tensorflow 2

Viewed 281

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?

0 Answers
Related