在 keras 中集成采样的 softmax 失败

基于如何在 Keras 模型中使用 TensorFlow 的采样 softmax 损失函数？ https://stackoverflow.com/questions/47892380/how-can-i-use-tensorflows-sampled-softmax-loss-function-in-a-keras-model，我创建了这段代码：

class SampledSoftmax(tensorflow.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(SampledSoftmax, self).__init__(**kwargs)


    def call(self, inputs):

        def f1(inputs):
            return tf.nn.sampled_softmax_loss(
                                inputs[0]._keras_history[0].weights[0],
                                inputs[0]._keras_history[0].bias,
                                tf.reshape(tf.argmax(inputs[1], 1), [-1, 1]),
                                inputs[0],                                
                                8192,
                                817496)
        def f2(inputs):
            logits = tf.matmul(inputs[0], tf.transpose(inputs[0]._keras_history[0].weights[0]))
            logits = tf.nn.bias_add(logits, inputs[0]._keras_history[0].bias)            
            return tf.nn.softmax_cross_entropy_with_logits_v2(
                                labels=inputs[1],
                                logits=logits)    

        return tf.cond(K.learning_phase(), true_fn=f1(inputs), false_fn=f2(inputs))

与以下型号一起使用时：

#model
input_layer = Input(shape=(None,), dtype='int32')
target_input = Input(shape=(None,vocab_size), dtype='int8')

embedding_layer = Embedding(vocab_size,
                            EMBEDDING_DIM,
                            trainable=True,
                            mask_zero=True) (input_layer)
common = LSTM(LSTM_UNITS, return_sequences=True,dropout=0.2, recurrent_dropout=0.2)(embedding_layer)
common = (Dense(PROJ_UNITS, activation='linear'))(common)
out = (Dense(vocab_size, name='output_layer'))(common)
out = (SampledSoftmax())([out, target_input])


model = Model(inputs=[input_layer,target_input], outputs=out)

它因以下错误而失败： ValueError：形状必须为等级 2，但对于“sampled_softmax/sampled_softmax_loss/MatMul”（操作：“MatMul”），形状必须为等级 3，输入形状为：[?,?,817496]、[?,817496]。

我根据谷歌搜索取得了一些进展：

class MyLayer(tensorflow.keras.layers.Dense):
    def __init__(self, num_sampled, num_classes, mode,  **kwargs):
        self.num_sampled = num_sampled
        self.num_classes = num_classes
        self.mode = mode
        super(MyLayer, self).__init__(num_classes, **kwargs)
        self.input_spec = [InputSpec(ndim=2)]

    def build(self, input_shape):
        #self.input_spec = [InputSpec(shape=input_shape)]
        super(MyLayer, self).build(input_shape)  # Be sure to call this somewhere!

    def call(self, inputs_and_labels):
        inputs, labels = inputs_and_labels
        if self.mode == "train":
            loss = tf.nn.sampled_softmax_loss(
                weights=self.kernel,
                biases=self.bias,
                labels=tf.reshape(tf.argmax(labels, 1), [-1, 1]),
                inputs=inputs,
                num_sampled=self.num_sampled,
                num_classes=self.num_classes,
                num_true=1)

        elif self.mode == "eval":
            logits = tf.matmul(inputs, tf.transpose(self.kernel))
            logits = tf.nn.bias_add(logits, self.bias)
            loss = tf.nn.softmax_cross_entropy_with_logits(
                labels=labels,
                logits=logits)

        return loss

    def compute_output_shape(self, input_shape):
        dense_shape, classes_shape = input_shape
        return (dense_shape[0], )    

现在的错误： 现在的错误：

ValueError: Layer my_layer expects 1 inputs, but it received 2 input tensors. Inputs received: [<tf.Tensor 'dense/BiasAdd:0' shape=(?, ?, 512) dtype=float32>, <tf.Tensor 'input_2:0' shape=(?, ?, 817496) dtype=int8>]

我尝试使用 self.input_spec 但直到现在才起作用。