Импорт Tensorflow 2.0 GPU из разных процессов

Question

Я работаю над проектом, в котором у меня есть модуль python, который реализует итеративный процесс, и некоторые вычисления выполняются с помощью графического процессора с использованием tennsflow 2.0. Модуль работает правильно, когда используется отдельно от одного процесса.

Поскольку мне нужно выполнить несколько прогонов с разными параметрами, я бы хотел распараллелить вызовы, но когда я вызываю модуль (который импортирует тензорный поток) из другого процесса, я получаю CUDA_ERROR_OUT_OF_MEMORY и бесконечный циклCUDA_ERROR_NOT_INITIALIZED, поэтому порожденные процессы зависают навсегда.

Конечно, я пытался ограничить память GPU , и это работает правильно, если я запускаю два разных скрипта Python от разных интерпретаторов, но, похоже,не работает в моем случае.

В частности, если я использую

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
       # Currently, memory growth needs to be the same across GPUs
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
            logical_gpus = tf.config.experimental.list_logical_devices('GPU')
            print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
    except RuntimeError as e:
        # Memory growth must be set before GPUs have been initialized
        print(e)

, я получаю бесконечный цикл CUDA_ERROR_NOT_INITIALIZED, а если я использую:

physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
    tf.config.experimental.set_memory_growth(physical_devices[0], True)
else:
    print("No GPU found, model running on CPU")

Процесс также зависает, но я получаю сообщение об ошибке для каждого порожденного процесса.

При чтении вывода консоли tenorflow кажется, что первый порожденный процесс выделяет память на GPU, однако он зависает так же, как и другие процессы. что жалуются на истощение памяти. Любопытно, что в nvidia-smi память графического процессора совсем не исчерпывается.

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 410.48                 Driver Version: 410.48                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  TITAN V             Off  | 00000000:03:00.0  On |                  N/A |
| 29%   42C    P8    28W / 250W |    755MiB / 12035MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

Мне удалось написать минимальный воспроизводимый пример проблемы:

File "tf_module. py "

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
       # Currently, memory growth needs to be the same across GPUs
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
            logical_gpus = tf.config.experimental.list_logical_devices('GPU')
            print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
    except RuntimeError as e:
        # Memory growth must be set before GPUs have been initialized
        print(e)
else:
    print("Running on CPU")

def run(x, y):
    return tf.add(x, y).numpy()

Файл" run.py "

from multiprocessing import Pool
import tf_module as experiment
def run_exp(params):
    a, b = params
    return experiment.run(a, b)

pool = Pool(2)
params = [(a, b) for a in range(3) for b in range(3)]

results = pool.map(run_exp, params)

Перенос вычисления TF из модуля невозможен, поскольку он является частью сложного конвейера, в котором также есть numpy. По этой причине мне нужно распараллелить этот фрагмент кода.

Я что-то упустил?

Заранее спасибо