Question

我正在设计一个GPU op内核，它在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));
}

将allocate_persistent（）与OpKernelConstruction上下文结合使用为积累提供持久的缓冲。我不想这样做，因为我处理的是可变缓冲区大小，它们可能相当大。

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运算符

1 个答案: