时间戳

首页 > 解决方案 > 将 Keras/TF/PyTorch 层分配给硬件类型

tensorflow - 将 Keras/TF/PyTorch 层分配给硬件类型

问题描述

假设我们有以下架构：

  1. 多个 CNN 层
  2. RNN层
  3. （时间分布）密集分类层

我们现在想训练这个架构。我们花哨的 GPU 在解决 CNN 层时非常快。虽然使用较低的时钟频率，但它可以并行执行许多卷积，从而提高速度。然而，我们花哨的 CPU 对于（非常长的）结果时间序列来说更快，因为时间步长不能并行化，并且处理从更高的 CPU 时钟速率中获益。所以执行的（据说）聪明的想法是这样的：

  1. 多个 CNN 层（在 GPU 上运行）
  2. RNN 层（在 CPU 上运行）
  3. （时间分布）密集分类层（在 GPU/CPU 上运行）

这让我想到了两个重要的问题：

使用标题中提到的任何框架，是否有可能将某些层分配到某些硬件，以及如何？

如果可能的话，额外内存操作的开销，例如 GPU-/CPU-RAM 之间的传输，是否会使整个想法变得毫无用处？

标签： tensorflowkerasdeep-learningpytorch

解决方案

基本上，在 Pytorch 中，您可以控制变量/参数所在的设备。AFAIK，您有责任确保对于每个操作，所有参数都驻留在同一设备上：即，您不能conv(x, y)在xGPUy上和 CPU 上。

这是通过 pytorch 的.to()方法完成的，该方法移动模块/变量.to('cpu')或.to('cuda:0')

推荐阅读