如何从 tf.math.segment_sum 创建 Keras 层答案

【问题标题】：How to create a Keras layer from tf.math.segment_sum如何从 tf.math.segment_sum 创建 Keras 层
【发布时间】：2022-06-22 16:49:54
【问题描述】：

我想在 Keras 层中使用 tf.math.segment_sum 函数，但尺寸不正确。

例如，我想在数据框df 中对x_1 分组的id 的值求和：

df = pd.DataFrame({'id':     [1, 1, 2, 2, 3, 3, 4, 4],
                   'x_1':    [1, 0, 0, 0, 0, 1, 1, 1],
                   'target': [1, 1, 0, 0, 1, 1, 2, 2]})

我创建的“模型”如下所示：

input_ = tf.keras.Input((1,), name='X')

cid = tf.keras.Input(shape=(1,), dtype='int64', name='id')

summed = tf.keras.layers.Lambda(lambda x: tf.math.segment_sum(x[0], x[1]), name='segment_sum')([input_, cid])

model = tf.keras.Model(inputs=[input_, cid], outputs=[summed])

我收到关于排名的错误：
ValueError: Shape must be rank 1 but is rank 2 for 'segment_sum/SegmentSum' (op: 'SegmentSum') with input shapes: [?,1], [?,1].

我在这里做错了什么？

【问题讨论】：

您是否尝试过扁平化参数x[0] 和x[1]？我认为删除维度（即tf.reshape(len(x[i]))）可能会解决问题。
我更改了以下行：summed = tf.keras.layers.Lambda(lambda x: tf.math.segment_sum(x[0], tf.reshape(x[1], (1,))), name='segment_sum')([input_, cid])。这消除了错误。当我做一个像model.predict([np.asarray([[6]]), np.asarray([0])]) 这样的“预测”时，我得到了预期的输出array([[6.]], dtype=float32)。但是，当我输入 model.predict([np.asarray([[6]]), np.asarray([5])]) 时，我希望得到相同的输出，但出现错误：ValueError: Mismatch between expected batch size and model output batch size. Output shape = (6, 1), expected output shape = shape (1, 1)。有什么想法吗？
你为什么要改成 (1,)？我想你想重塑为 (len(x[1]),)，正如我上面所说的。
len() 不适用于 tensorflow 张量...您知道另一种方法吗？

标签： tensorflow keras

【解决方案1】：

我使用tf.gather 解决了这个问题。工作代码如下：

input_ = tf.keras.Input((1,), name='X')

cid = tf.keras.Input(shape=(1,), dtype='int64', name='id')

summed = tf.keras.layers.Lambda(lambda x: tf.gather(tf.math.segment_sum(x[0], tf.reshape(x[1], (-1,))), x[1]), output_shape=(None,1), name='segment_sum')([input_, cid])

model = tf.keras.Model(inputs=[input_, cid], outputs=[summed])

【讨论】：