支持 Tensorflow 2 在单 GPU 上进行多处理推理

Question

我有一组图像需要进行推理。我想做的是生成 4 个工作线程，每个工作线程对 GPU 内存有 1/4 的访问权限。即使我将每个工作人员的内存限制为 1/8，我也会遇到 OOM。下面是一些代码：

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)

在

do_predict()

函数中，我导入tensorflow并尝试了这两个内存分配选项：

# 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)

两者都给我 OOM 错误。根据 nvidia-smi，我的 GPU 有 8081 MiB 可用空间，因此为每个进程分配 1024 MB 应该没问题。所以我一定是在以错误的方式做某事或思考某事。有什么想法吗？

Answer 1

虽然您可能没有使用卡提供给您的所有内存，但也许您的进程正在耗尽分配给它们的所有内存。换句话说，也许给他们更少的记忆实际上是相反的方向？