tensorflow tf.slice会产生分配和/或内存复制吗?

时间:2018-06-10 01:28:33

标签: python tensorflow memory-management

确实

b = tf.slice(a, [...], [...])

分配一个新的内存缓冲区然后从缓冲区复制?或者a和b共享同一个缓冲区?

那么,

... = tf.nn.convolution(tf.slice(a, [...], [...]), ...)

在这种情况下,切片未被命名。这会产生分配和复制吗?

一般来说,是否有一些资源可用于学习这些内存管理细节?

谢谢!

1 个答案:

答案 0 :(得分:0)

我也对此感到很好奇,并研究了tf.slice的内核实现。

唯一的一次不分配内存的时间是当切片是身份切片(因此没有更改)时,以及切片是0维对齐的。

Identity Lines

  if (is_identity) {
    VLOG(1) << "Slice identity";
    context->set_output(0, input);
    *done = true;
    return;
  }

Dim0Aligned

  if (slice_dim0 &&
      IsDim0SliceAligned<T>(input.shape(), (*begin)[0], (*size)[0])) {
    VLOG(1) << "Slice dim 0: " << input.shape().DebugString();
    CHECK_GE(input.dims(), 1);  // Otherwise, is_identity should be true.
    context->set_output(0, input.Slice((*begin)[0], (*begin)[0] + (*size)[0]));
    *done = true;
    return;
  }

否则,以下将被称为

  OP_REQUIRES_OK(context, context->allocate_output(0, *output_shape, result));