使TensorFlow使用通过自定义CUDA例程即时生成的训练数据

Question

假设我们生成了自己的训练数据（例如，从某个扩散过程中采样并计算了一些感兴趣的量），并且我们拥有自己的CUDA例程，称为generate_data，该例程在GPU内存中为给定的一组输入生成标签。

因此，我们处于一种特殊的设置中，可以“在线”方式生成任意数量的训练数据（在每个批次迭代中，我们都调用generate_data例程以生成新批次并丢弃旧批次）。

由于数据是在GPU上生成的，有没有办法使TensorFlow（Python API）在训练过程中直接使用它？（例如填充一个占位符）这样，这样的管道将是有效的。

我的理解是，当前您需要在这样的设置中将数据从GPU复制到CPU，然后让TensorFlow将数据再次从CPU复制到GPU，这在执行不必要的复制时非常浪费。

编辑：如果有帮助，我们可以假定CUDA例程是使用Numba的CUDA JIT编译器实现的。

Answer 1

这绝对不是一个完整的答案，但希望能有所帮助。

• 您可以通过编写自定义操作将 CUDA 例程集成到 TensorFlow 。目前 TensorFlow 中没有其他方式与其他 CUDA 例程交互。

• 至于完全在 GPU 上编写训练循环，我们可以使用与此SO 问题tf.while_loop非常相似的方式在 GPU 上编写例程：

  i = tf.Variable(0, name='loop_i')
  
  def cond(i):
      return i < n
  
  def body(i):
      # Building the graph for custom routine and our model
      x, ground_truth = CustomCUDARountine(random_seed, ...)
      predictions = MyModel(x, ...)
  
      # Defining the optimizer
      loss = loss_func(ground_truth, predictions)
      optim = tf.train.GradientDescentOptimizer().minimize(loss)
  
      # loop body
      return tf.tuple([tf.add(i, 1)], control_inputs=[optim])
  
  loop = tf.while_loop(cond, body, [i])
  
  # Run the loop
  tf.get_default_session().run(loop)