时间戳

首页 > 解决方案 > tf.test.is_gpu_available() 在子进程中为 False，但在主进程中为 True

tensorflow - tf.test.is_gpu_available() 在子进程中为 False，但在主进程中为 True

问题描述

我目前正在运行一个 pytorch 模型，该模型会定期调用 tensorflow 模型以进行基准测试。我希望这两个模型都支持 GPU 并在同一个脚本中运行。由于 tensorflow 基准测试代码要求 GPU 内存直到流程结束，我选择在运行基准测试代码时运行基准测试代码，multiprocessing.Process以便我的 pytorch 模型可以在基准测试脚本运行后使用完整的 GPU 内存。

在此期间，我在 tensorflow 的 gpu 利用率中偶然发现了一个不寻常的错误（？）。似乎在子进程中运行的 tensorflow 不想使用已被父进程使用的 GPU。我可以在同一个 GPU 中拥有 tensorflow 模型和 pytorch 模型，并且进程没有问题，但是当我引入子进程时，tensorflow 表现不佳。

我 tensorflow-gpu==1.14.0 torch==1.1.0 cudatoolkit=10.0在 NVIDIA 2080-Ti 上运行。

以下是截取的最小代码以重现：

import torch
import tensorflow as tf
from multiprocessing import Process

def f():
    print(tf.test.is_gpu_available())

pa = Process(target=f, args=())
pa.start()
pa.join()

torch.ones(1).cuda()

pb = Process(target=f, args=())
pb.start()
pb.join()

>>> True
>>> False

标签： tensorflowmultiprocessingpytorch

解决方案

对于遇到此问题的任何人，您需要致电multiprocessing.set_start_method('spawn'). Tensorflow 不是分叉安全的，全局变量/模块可能会发生一些奇怪的事情，这可能很难推理。请记住只在支票内调用一次if __name__ == '__main__':。

推荐阅读