在自定义层中将可变大小张量填充到特定大小

Question

我有一个二维张量，第一维是动态的，第二维是固定的。我使用segment_sum更新这个张量，其中第一个维度的大小可能会改变，因此，我想将修改后的张量零填充到与输入相同的形状。

在这个问题中提供的答案对我没有帮助，因此用我的用例的细节来问这个问题。

class MyLayer(layers.Layer):
def call(self, inputs):
x, segment_ids = inputs
x_ = tf.math.segment_sum(x, segment_ids)
# the first dimension of x and x_ 
# may not be equal at this point, hence zero-padding.
x_ = tf.pad(
x_, 
[(0, 0), (tf.shape(x)[0] - tf.shape(x_)[0], 0)])
return x_

如果我正在阅读如何填充未知来固定大小并正确获得动态大小张量的形状，上述应该工作!?然而，我得到一个错误在层下游MyLayer抱怨:

不支持*=:'int'和'NoneType'的操作数类型

我得到这个错误，甚至在第一个epoch开始之前，不管我提供的输入;然而，如果我删除填充，我将得到不兼容的形状只有我的一些输入。因此，我猜测错误与填充有关，也许tf.shape在错误消息中给出NoneType?

---

这是一个最小的示例，我可以使用简单的填充来获得预期的输出，如示例末尾所示。然而，在我的实际用例中，模型是根据动态输入大小定义的，即，大小为None，类似的方法会失败，导致上述错误。

x = [[1, 10], [2, 20], [3, 30], [4, 40]]
i = [0, 1, 1, 1]
o = tf.math.segment_sum(x, i)
o = tf.get_static_value(o)
print(f"shape: {o.shape}")
print(f"ntensor:n{o}")

输出:

shape: (2, 2)
tensor:
[[ 1 10]
[ 9 90]]

预期输出

o = tf.pad(o, [(0, len(x) - o.shape[0]), (0, 0)])
shape: (4, 2)
tensor:
[[ 1 10]
[ 9 90]
[ 0  0]
[ 0  0]]

Answer 1

您可以尝试使用tf.concat在您的情况下填充:

import tensorflow as tf
class MyLayer(tf.keras.layers.Layer):
def call(self, inputs):
x, segment_ids = inputs
x_ = tf.math.segment_sum(x, segment_ids)
x_ = tf.concat([x_, tf.zeros((int(tf.shape(x)[0] - tf.shape(x_)[0]), tf.shape(x)[1]))], axis=0)
return x_
inputs = tf.keras.layers.Input((2,))
l = MyLayer()
x = l([inputs, [0, 1, 1, 1]])
outputs = tf.keras.layers.Dense(1)(x)
model = tf.keras.Model(inputs, outputs)
model.compile(optimizer='adam', loss='mse')
x = tf.constant([[1, 10], [2, 20], [3, 30], [4, 40]], dtype=tf.float32)
print(l([x, [0, 1, 1, 1]]))
y = tf.constant([1, 2, 1, 4])
model.fit(x, y, epochs=2)

tf.Tensor(
[[ 1. 10.]
[ 9. 90.]
[ 0.  0.]
[ 0.  0.]], shape=(4, 2), dtype=float32)
Epoch 1/2
1/1 [==============================] - 1s 683ms/step - loss: 2624.1050
Epoch 2/2
1/1 [==============================] - 0s 12ms/step - loss: 2361.9163
<keras.callbacks.History at 0x7f7fdd4cf810>