tensorflow2.0 squeeze出错

Posted 2022-11-30 lolybj

用tf.keras写了自定义层，但在调用自定义层的时候总是报错，找了好久才发现问题所在，所以记下此问题。

问题代码

u=tf.squeeze(tf.expand_dims(tf.expand_dims(inputs,axis=1),axis=3)@self.kernel,axis=3)

其中inputs的第一维为None，这里的代码为自定义的前向传播。我是想将得到的输出张量维度为1的维度删掉，因此调用了tf.squeeze方法，这时虽然没有报错但出现了问题。我分别打印了下面内容。

print(tf.expand_dims(tf.expand_dims(inputs,axis=1),axis=3).shape)
print(self.kernel.shape)
print((tf.expand_dims(tf.expand_dims(inputs,axis=1),axis=3)@self.kernel).shape)
print(tf.squeeze(tf.expand_dims(tf.expand_dims(inputs,axis=1))@self.kernel,axis=3))

可以发现，当张量第一维为None的时候tf.squeeze使结果变为了0。我想要的结果是删除第三个输出的大小为1的维度，即得到下面的结果

解决使用tf.squeeze的时候加上删除的维度。

tf.squeeze(tf.expand_dims(tf.expand_dims(inputs,axis=1),axis=3)@self.kernel,axis=3)