TensorFlow中实现NumPy np.resize等效功能的方法

Great question! The np.resize function has a specific, handy behavior: it flattens the input array, repeats it enough times to hit the total element count required by the target shape, truncates any excess elements, then reshapes the result. Let's build this exact logic using pure TensorFlow operations.

Step-by-Step Implementation #

Here's a function that mirrors np.resize perfectly:

import tensorflow as tf

def tf_resize(x, shape):
    # Calculate total elements needed for the target shape
    target_total = tf.reduce_prod(tf.convert_to_tensor(shape))
    # Flatten the input tensor to 1D
    x_flat = tf.reshape(x, (-1,))
    x_length = tf.shape(x_flat)[0]
    
    # Calculate how many times to repeat the flattened tensor (ceil division)
    repeat_count = tf.cast(tf.math.ceil(target_total / tf.cast(x_length, tf.float32)), tf.int32)
    # Repeat the tensor to get enough elements
    x_repeated = tf.tile(x_flat, [repeat_count])
    # Truncate to exactly the number of elements we need
    x_truncated = x_repeated[:target_total]
    # Reshape to the target shape
    return tf.reshape(x_truncated, shape)

Test with Your Examples #

Let's verify this works with the cases you provided:

Example 1: Truncating excess elements

x = tf.constant([1, 2, 3, 4, 5, 6])
y = tf_resize(x, shape=(2, 2))
# Confirm the result matches numpy's output
assert tf.equal(y, tf.constant([[1, 2], [3, 4]])).numpy().all()

Example 2: Repeating to fill missing elements

z = tf_resize(x, shape=(3, 3))
# Confirm the result matches numpy's output
assert tf.equal(z, tf.constant([[1, 2, 3], [4, 5, 6], [1, 2, 3]])).numpy().all()

Quick Notes #

• This works with any input tensor shape (not just 1D)—just like np.resize, it flattens the input first no matter its original dimensions.
• We use tf.tile for repeating elements, which is TensorFlow's native way to duplicate tensor segments along an axis.
• Ceil division ensures we always generate enough elements before truncating to the exact target size.

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