如何在单个程序中运行运行 2 个单独的 tf.keras 模型，一个在 CPU 上，另一个在 GPU 上

Question

我正在训练一个深度神经网络模型，-让我们称之为模型 A- 在我的主要训练脚本 train.py 上使用 tf.keras。此 train.py 导入 utils.py，其中定义了用于训练过程的数据生成器。 utils.py 还导入了 other_utils.py，其中创建并初始化了预训练的深度神经网络模型 -MODEL B-，因为数据生成器使用它来处理训练数据。

模型 B 的初始化已经使用了大部分 GPU 内存，让我缺乏资源来训练模型 A。

导入模式如下所示：

train.py (MODEL A) --> utils.py (data generator) --> other_utils.py (MODEL B)

所以我想运行模型B在CPU上生成数据，然后在GPU上训练模型A。

为了防止MODEL B使用GPU，我尝试在other_上设置环境变量_utils.py

os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

然后将其设置回 utils.py 和 train.py 然后导入 tensorflow

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import tensorflow

虽然设置 CUDA_VISIBLE_DEVICES = -1 会禁用 MODEL B 的 GPU，但它也会禁用 GPU 以供进一步使用并在 train.py 上将其设置回 CUDA_VISIBLE_DEVICES = 0，然后导入 tensorflow不会启用它。

那么如何在 CPU 上使用模型 B 并在 GPU 上使用模型 A？

Answer 1

我使用解决方案 here 来实现我的程序。

在 CPU 上执行：

with tf.device('/cpu:0'):
    a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
    b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
    c = tf.matmul(a, b)

    print(c.device)

输出：

/job:localhost/replica:0/task:0/device:CPU:0

在 GPU 上执行：

with tf.device('/gpu:0'):
    a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
    b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
    c = tf.matmul(a, b)

    print(c.device)

输出：

/job:localhost/replica:0/task:0/device:GPU:0

虽然，对于任何将在数据生成器中使用预训练模型来训练另一个模型的人来说，我必须说阻止预训练模型使用 GPU 并没有为我提供额外的 GPU 训练加速。

如何在单个程序中运行运行 2 个单独的 tf.keras 模型，一个在 CPU 上，另一个在 GPU 上

How tun run 2 seperate tf.keras models in a single program one on CPU and another on GPU

python

deep-learning

keras

tensorflow

如何在单个程序中运行 运行 2 个单独的 tf.keras 模型，一个在 CPU 上，另一个在 GPU 上

How tun run 2 seperate tf.keras models in a single program one on CPU and another on GPU

python

deep-learning

keras

tensorflow

如何在单个程序中运行运行 2 个单独的 tf.keras 模型，一个在 CPU 上，另一个在 GPU 上