TensorFlow中如何对特定张量内x、y向量求和以调整维度？

Hey there! Let's tackle this problem step by step. Your goal is to take a tensor G with shape [batch_size, C, 1 + x + y, feature_size] and reduce its third dimension (0-based index 2) into 3 segments: the first element, the next x elements, and the final y elements—summing each segment to get a new tensor of shape [batch_size, C, 3, feature_size].

Here's a clean, efficient way to do this using core TensorFlow operations:

Step-by-Step Implementation #

1. Split the tensor into three segments: Use tf.split to divide the third dimension into parts of size 1, x, and y respectively.
2. Sum each segment: For each split part, apply tf.reduce_sum along the third dimension, keeping the dimension intact with keepdims=True (this ensures we don't lose the axis needed for concatenation).
3. Concatenate the summed segments: Combine the three summed tensors back along the third dimension to get your target shape.

Code Example

import tensorflow as tf

# Assume G is your input tensor with shape [batch_size, C, 1+x+y, feature_size]
# Define x and y (these can be static integers or dynamic tensors)
x = ...  # Your x value
y = ...  # Your y value

# Split the tensor into 3 parts: [1, x, y] along axis 2
split_parts = tf.split(G, num_or_size_splits=[1, x, y], axis=2)

# Sum each part along axis 2, keeping the dimension
summed_parts = [tf.reduce_sum(part, axis=2, keepdims=True) for part in split_parts]

# Concatenate the summed parts along axis 2 to get shape [batch_size, C, 3, feature_size]
new_G = tf.concat(summed_parts, axis=2)

Why This Works #

• tf.split handles both static and dynamic sizes for x and y, so it works even if these values aren't known at graph construction time.
• Using keepdims=True ensures each summed segment maintains a shape of [batch_size, C, 1, feature_size], making concatenation straightforward.
• The final tf.concat merges the three 1-length segments into a single 3-length dimension, exactly matching your desired output shape.

If you want to optimize further (e.g., avoid splitting), you could also use slicing directly:

# Slice each segment and sum
part1 = G[:, :, 0:1, :]  # Sum isn't needed here since it's a single element
part2 = tf.reduce_sum(G[:, :, 1:1+x, :], axis=2, keepdims=True)
part3 = tf.reduce_sum(G[:, :, 1+x:, :], axis=2, keepdims=True)

new_G = tf.concat([part1, part2, part3], axis=2)

This achieves the same result and might be slightly faster for very large tensors, but the split-based approach is more readable.

内容的提问来源于stack exchange，提问作者save_ole