TensorFlow：如何以-1为分隔符分割指定一维张量？

没问题，我来教你怎么实现这个需求！在TensorFlow里，要把给定的张量按-1分割成多个子张量，我们可以通过定位分隔符位置、计算各段的起止索引，再提取对应元素来完成。下面是两种实用的实现方式：

方法一：Eager模式下的直观实现 #

这种方式用Python循环配合TensorFlow的基础操作，逻辑清晰，适合日常调试和Eager模式下的场景：

import tensorflow as tf

# 定义原始张量
ts = tf.constant([1,2,3,-1,3,4,5,-1,2])

# 1. 找到所有-1的位置索引，转成一维张量
split_indices = tf.squeeze(tf.where(ts == -1), axis=1)

# 2. 构造每个子张量的起始和结束位置
# 起始点：从0开始，之后每个分隔符的下一个位置
starts = tf.concat([[0], split_indices + 1], axis=0)
# 结束点：每个分隔符的位置，最后以张量末尾结束
ends = tf.concat([split_indices, [tf.shape(ts)[0]]], axis=0)

# 3. 遍历起止索引，提取子张量并收集到列表中
split_tensors = []
for start, end in zip(starts.numpy(), ends.numpy()):
    # 过滤掉可能的空段（比如张量开头/结尾是-1的情况）
    if start < end:
        split_tensors.append(tf.constant(ts[start:end]))

# 输出结果
for idx, tensor in enumerate(split_tensors):
    print(f"第{idx+1}个子张量：{tensor}")

运行后你会得到：

第1个子张量：tf.Tensor([1 2 3], shape=(3,), dtype=int32)
第2个子张量：tf.Tensor([3 4 5], shape=(3,), dtype=int32)
第3个子张量：tf.Tensor([2], shape=(1,), dtype=int32)

方法二：纯TensorFlow图模式实现 #

如果你的代码需要在计算图模式下运行（比如用于模型部署），可以用tf.map_fn替代Python循环，全程用TensorFlow操作：

import tensorflow as tf

ts = tf.constant([1,2,3,-1,3,4,5,-1,2])

# 同样先定位分隔符索引
split_indices = tf.squeeze(tf.where(ts == -1), axis=1)
starts = tf.concat([[0], split_indices + 1], axis=0)
ends = tf.concat([split_indices, [tf.shape(ts)[0]]], axis=0)

# 把起止索引配对成二维张量
index_pairs = tf.stack([starts, ends], axis=1)

# 定义提取子张量的函数
def extract_segment(pair):
    start, end = pair[0], pair[1]
    return ts[start:end]

# 用tf.map_fn批量提取
split_tensors = tf.map_fn(extract_segment, index_pairs, dtype=tf.int32)

# 过滤空张量（可选，根据你的输入情况调整）
split_tensors = [t for t in split_tensors if tf.shape(t)[0] > 0]

print(split_tensors)

这段代码的输出和方法一完全一致，而且可以无缝融入TensorFlow的计算图中。

小提示 #

如果你的输入张量可能存在连续的-1（比如[1,-1,-1,2]），上面的代码会自动处理，因为每个-1都会被当作分隔符，中间的空段会被过滤掉，最终得到[tf.constant([1]), tf.constant([2])]。

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