Python TensorFlow：如何自定义收集/拆分张量？

请允许我简单介绍一下背景。我正在研究Bi LSTM CRF体系结构，以完成序列标记任务。我不想将LSTM编码的各个步骤的特征输入CRF层，而是想收集属于同一段的特征向量，并对它们进行平均，然后将段特征发送到CRF层。该方法是在序列组块的神经模型中提出的

现在我们有一个3D张量

X

[batch\u size，num\u steps，num\u features]和一个2D张量

Xl

[batch\u size，num\u steps]，指示序列的IOB标签。我如何收集属于同一段的特征向量，并将其转换为4D输出张量[batch_size，max_segment_num，max_segment_len，num_features]？段边界由IOB标签确定，其中B表示开始，I表示内部标记

X:
[[v11, v12, v13, v14, v15],
 [v21, v22, v23, v24, v25],
 ...]

IOB labels:
[[B, I, B, I, I],
 [B, B, I, I, B],
 ...]

and the expected output:
[[[v11, v12], [v13, v14, v15]],
 [[v21], [v22, v23, v24], [v25]],
 ...]

with the output, we can compute segment features:
[[u11, u12, 0],  # supposing max_segment_num = 3
 [u21, u22, u23],
 ...]

我研究了一些内置函数，如

tf.gather

或

tf.split

，但据我所知，它们都不支持自定义方式。我也研究了掩蔽函数，但没有找到它们能起到什么作用

我试图手动操纵张量并分配给新创建的4D张量，但被卡住了。模型定义中的张量X是一个占位符，在训练中输入数据。因此，下面的语句失败了，抱怨我们现在必须将数据提供给占位符

for i in range(batch_size):
    ...
    for j in range(max_segment_num):
        ...
        seg_tensor = tf.gather(X[i,:], rng)  # rng is the segment range
        average_feature = tf.reduce_mean(seg_tensor, axis=0)
        update_op = outputs[i, j, :].assign(average_feature)
        update_op.run()

谁能给我一个提示吗？我如何解决这个问题，或者是否有更简单的方法使用内置函数？提前感谢。

您可以使用它使numpy数组上的任何计算行为类似于本机TensorFlow op。