Skip to content Skip to sidebar Skip to footer

Does Tensorflow Tf.slice Incur Allocation And/or Memory Copy?

Does b = tf.slice(a, [...], [...]) allocate a new memory buffer and then copy from a's buffer? or do a and b share the same buffer? And, what about ... = tf.nn.convolution(tf.sli

Solution 1:

I was curious of this as well and went into the kernel implementation of tf.slice.

The only time memory is not allocated is when the slice is the identity slice (so there are no changes), and when the slice is dimension 0 aligned.

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;
  }

Otherwise the following will be called

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

Post a Comment for "Does Tensorflow Tf.slice Incur Allocation And/or Memory Copy?"