在 tensorflow 中实现软 argmax

Question

我正在尝试在 Tensorflow 中实现软 argmax 操作。 我想在前向传递中使用正常的 argmax，在后向传递中使用 softmax 近似值。 准确地说，我的输入是 NCHW 格式，其中 C 是 2 个通道。 我现在面临的问题是由于切片而导致的降维。 我的渐变和 output 具有相同的形状（1通道输出），但是据我了解，它们与 Tensorflow 不允许的输入不同。

tensorflow.python.framework.errors_impl.InvalidArgumentError: Input to reshape is a tensor with 409600 values, but the requested shape has 819200

我目前最好的尝试：

@tf.custom_gradient
def soft_argmax(x):
    out_no_grad = tf.argmax(x, axis=1)

    @tf.function
    def argmax_soft(x):
        out = tf.nn.softmax(x, axis=1)
        return

    def grad(dy):
        gradient = tf.gradients(argmax_soft(x), x)[0][:,0,:,:]
        return gradient * dy
return out_no_grad, grad

我怎么能设计出这样的 function？ 请注意，我专门搜索定义自定义渐变或 function 的东西。 我知道我可以做grad + tf.stop_gradient(out_no_grad - grad) ，但这不是我想要的。