python - Keras 中 CPU 和 GPU 的混合使用

Question

我正在 Keras 上构建一个神经网络，包括多层 LSTM、Permute 和 Dense。

似乎 LSTM 对 GPU 不友好。所以我做了研究和使用

With tf.device('/cpu:0'):
   out = LSTM(cells)(inp)

但是根据我对with的理解，with是try...finally block ，以确保执行清理代码。我不知道下面的 CPU/GPU 混合使用代码是否有效？他们会加快训练速度吗？

With tf.device('/cpu:0'):
  out = LSTM(cells)(inp)
With tf.device('/gpu:0'):
  out = Permute(some_shape)(out)
With tf.device('/cpu:0'):
  out = LSTM(cells)(out)
With tf.device('/gpu:0'):
  out = Dense(output_size)(out)

Answer 1

如您所见here - tf.device 是一个上下文管理器，它将默认设备切换到这个作为它创建的上下文( block )中的参数传递的设备。所以这段代码应该在 CPU 上运行所有 '/cpu:0' 设备并在 GPU 上休息。

它会加速你的训练这个问题真的很难回答，因为它取决于你使用的机器 - 但我不希望计算速度更快，因为每次更换设备都会在 之间复制数据GPU RAM 和机器 RAM。这甚至可能会减慢您的计算速度。