设计一个累积的 Tensorflow GPU 算子

Question

我正在设计一个 GPU 运算内核，它在 GPU 内存的缓冲区中迭代地累积数据。数据保留在 GPU 内存中很重要。所以类似的东西：

with tf.device('/gpu:0'):
    buffer = tf.zeros(...)

    buffer = accumulate(param11, param12, buffer)
    buffer = accumulate(param21, param22, buffer)
    buffer = accumulate(param31, param32, buffer)

with tf.device('/cpu:0'):
    A = do_some_more_stuff(buffer)

我想就我认为可用于实现此目的的三种方法提供一些意见：

在每次调用时分配输出张量并将其用作输入张量在下一次通话中。这很容易实现，但我担心 GPU内存的持续分配将是一个问题。 tensorflow 会将现在未使用的分配释放到 GPU 内存池中吗？

REGISTER_OP("Accumulate")
    .Input("param1: T")
    .Input("param2: T")
    .Input("buffer_in: T")
    .Output("buffer_out: T")

void Compute(tensorflow::OpKernelContext * ctx) override
{
    TensorShape output_shape{...};
    Tensor * output_ptr = nullptr;
    OP_REQUIRES_OK(ctx, ctx->allocate_output(
        0, output_shape, &output_ptr))

    kernel<<<grid, blocks, 0, stream>>>(
        ctx->input(0), ctx->input(1),
        output);
}

参考输入和输出张量并确保它们是参考到相同的数据。据我了解标准操作和 OpKernelContext 文档，这需要像其他操作一样用互斥锁保护也可能正在访问底层引用张量...

REGISTER_OP("Accumulate")
    .Input("param1: T")
    .Input("param2: T")
    .Input("buffer_in: Ref(T)")
    .Output("buffer_out: Ref(T)")

void Compute(tensorflow::OpKernelContext * ctx) override
{
    mutex_lock(mu_);

    ctx->forward_ref_input_to_ref_output(2, 0);

    kernel<<<grid, blocks, 0, stream>>>(
        ctx->input(0), ctx->input(1),
        ctx->mutable_input(2, true));
}

结合 OpKernelConstruction 上下文使用 allocate_persistent() 为积累提供持久缓冲。我不想这样做，因为我正在处理可变缓冲区大小，它们可能会相当大。

Answer 1

我不太确定你想用你的 C++ 代码做什么，但通过查看 python 代码片段，我认为 tf.assign 可能会有所帮助。它允许你做这样的事情：

buffer = tf.Variable(...)
param = tf.Variable(...)
accumulate_op = buffer.assign(expr<param, buffer>)

...

sess.run(accumulate_op)

运行 accumulate_op 应该更新 gpu 上的缓冲区（您可能必须将其包装在 tf.group 中以避免获取更新的值）。

设计一个累积的 Tensorflow GPU 算子

Designing an accumulating Tensorflow GPU operator

tensorflow