TensorFlow中如何拆分未知尺寸张量为指定大小块？求类似np.array_split的方法

Great question! I’ve struggled with this exact problem when working with tensors where the runtime shape isn’t known upfront—tf.split’s strict requirement for equal splits is definitely a pain point when you can’t guarantee even division.

Unfortunately, TensorFlow doesn’t have a direct equivalent to NumPy’s np.array_split as of the latest 2.x versions. But don’t worry—we can easily replicate that behavior with a custom function that dynamically calculates split sizes based on the tensor’s runtime shape.

方法1：自定义实现类似np.array_split的函数 #

This is the most flexible approach, supporting both splitting into a fixed number of chunks (with the last few potentially smaller) and splitting by specified chunk sizes (adjusting the final chunk to fit the remaining elements).

Here’s a robust implementation:

import tensorflow as tf

def tf_array_split(tensor, num_or_size_splits, axis=0):
    # 获取张量在目标轴上的运行时尺寸
    tensor_axis_size = tf.shape(tensor)[axis]
    
    if isinstance(num_or_size_splits, int):
        # 场景1：拆分为N个块（对应np.array_split(arr, N)）
        base_split_size = tensor_axis_size // num_or_size_splits
        remainder = tensor_axis_size % num_or_size_splits
        
        # 生成拆分尺寸：前remainder个块多1个元素，其余块为基础尺寸
        split_sizes = tf.concat([
            tf.fill([remainder], base_split_size + 1),
            tf.fill([num_or_size_splits - remainder], base_split_size)
        ], axis=0)
    else:
        # 场景2：按指定尺寸拆分（对应np.array_split(arr, [a, b, c])）
        cumulative_sizes = tf.cumsum(num_or_size_splits)
        # 调整最后一个块的尺寸以适配剩余元素
        final_chunk_size = tensor_axis_size - cumulative_sizes[-2] if tf.size(cumulative_sizes) > 1 else tensor_axis_size
        split_sizes = tf.concat([num_or_size_splits[:-1], [final_chunk_size]], axis=0)
    
    # 使用动态计算的尺寸执行拆分
    return tf.split(tensor, split_sizes, axis=axis)

用法示例 #

# 示例1：拆分为3个块（张量长度无法被3整除时，块尺寸不等）
tensor = tf.range(7)  # 形状 [7]
splits = tf_array_split(tensor, 3)
# 输出: [<tf.Tensor: shape=(3,), dtype=int32, numpy=array([0, 1, 2])>,
#          <tf.Tensor: shape=(2,), dtype=int32, numpy=array([3, 4])>,
#          <tf.Tensor: shape=(2,), dtype=int32, numpy=array([5, 6])>]

# 示例2：按指定尺寸拆分（最后一块自动调整以适配剩余元素）
tensor = tf.range(10)
splits = tf_array_split(tensor, [3, 4], axis=0)
# 输出: [<tf.Tensor: shape=(3,), dtype=int32, numpy=array([0, 1, 2])>,
#          <tf.Tensor: shape=(4,), dtype=int32, numpy=array([3, 4, 5, 6])>,
#          <tf.Tensor: shape=(3,), dtype=int32, numpy=array([7, 8, 9])>]

方法2：利用tf.data.Dataset（适用于一维张量） #

如果你处理的是一维张量，且偏好使用TensorFlow的数据集API，可以通过创建指定尺寸的窗口来拆分：

def split_via_dataset(tensor, num_splits, axis=0):
    if axis != 0:
        raise ValueError("This method only supports axis=0 for simplicity")
    
    tensor_size = tf.shape(tensor)[0]
    base_size = tensor_size // num_splits
    remainder = tensor_size % num_splits
    
    dataset = tf.data.Dataset.from_tensor_slices(tensor)
    splits = []
    
    # 先处理尺寸较大的块
    for _ in range(remainder):
        splits.append(tf.stack(list(dataset.take(base_size + 1))))
        dataset = dataset.skip(base_size + 1)
    
    # 处理剩余尺寸相等的块
    for _ in range(num_splits - remainder):
        splits.append(tf.stack(list(dataset.take(base_size))))
        dataset = dataset.skip(base_size)
    
    return splits

关键注意事项 #

• 处理动态形状时，务必使用tf.shape(tensor)而非tensor.shape——tensor.shape返回的是静态形状信息，对于运行时尺寸未知的张量可能返回None。
• 自定义的tf_array_split函数支持多维张量，只需指定axis参数即可沿目标维度拆分，和np.array_split行为完全一致。

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