如何在Keras模型中使用TensorFlow的采样softmax损失函数？

Question

我正在Keras训练语言模型,并希望通过使用采样softmax作为我网络中的最终激活功能来加速训练.从TF文档中,看起来我需要为weights和提供参数biases,但我不确定这些的输入是什么.好像我可以在Keras中编写自定义函数,如下所示:

import keras.backend as K

def sampled_softmax(weights, biases, y_true, y_pred, num_sampled, num_classes):
    return K.sampled_softmax(weights, biases, y_true, y_pred, num_sampled, num_classes)

但是,我不确定如何"插入"到我现有的网络.LM的架构非常简单:

model = Sequential()
model.add(Embedding(input_dim=len(vocab), output_dim=256))
model.add(LSTM(1024, return_sequence=True))
model.add(Dense(output_dim=len(vocab), activation='softmax'))
model.compile(loss='categorical_crossentropy', optimizer='adam')

鉴于这种架构,我可以在模型上调用编译方法时将sampled_softmax函数作为loss参数传递吗？或者这需要写为最终完全连接层之后的层.这里的任何指导将不胜感激.谢谢.

Answer 1

这里的关键观察是 TensorFlow 采样的 softmax 函数返回实际损失，而不是对可能标签集的一组预测，以与地面实况数据进行比较，然后作为单独的步骤计算损失。这使得模型设置有点奇怪。

首先，我们向模型添加第二个输入层，除了作为目标输出之外，该层第二次将目标（训练）数据编码为输入。这用于函数的labels参数sampled_softmax_loss。它必须是 Keras 输入，因为当我们去实例化和设置模型时，它被视为输入。

其次，我们构建了一个新的自定义 Keras 层，该层调用sampled_softmax_loss具有两个 Keras 层作为其输入的函数：预测我们类别的密集层的输出，然后是包含训练数据副本的第二个输入。请注意，我们正在访问_keras_history实例变量以从原始全连接层的输出张量中获取权重和偏置张量。

最后，我们必须构建一个新的“哑”损失函数，它忽略训练数据而只使用sampled_softmax_loss函数报告的损失。

请注意，由于采样的 softmax 函数返回的是损失，而不是类别预测，因此您不能使用此模型规范进行验证或推理。您需要在新规范中重新使用此“训练版本”中的训练层，该规范将标准 softmax 函数应用于应用了默认激活函数的原始密集层。

肯定有一种更优雅的方法来做到这一点，但我相信这是有效的，所以我想我现在就按原样张贴在这里，而不是等到我有一些更整洁的东西。例如，您可能希望将类的数量作为SampledSoftmax层的参数，或者更好的是，将所有这些都压缩到原始问题中的损失函数中，并避免两次传入训练数据。

from keras.models import Model
from keras.layers import Input, Dense, Layer
from keras import backend as K

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


    def call(self, inputs):
        """
        The first input should be the model as it were, and the second the
        target (i.e., a repeat of the training data) to compute the labels
        argument

        """
        # the labels input to this function is batch size by 1, where the
        # value at position (i, 1) is the index that is true (not zero)
        # e.g., (0, 0, 1) => (2) or (0, 1, 0, 0) => (1)
        return K.tf.nn.sampled_softmax_loss(weights=inputs[0]._keras_history[0].weights[0],
                                            biases=inputs[0]._keras_history[0].bias,
                                            inputs=inputs[0],
                                            labels=K.tf.reshape(K.tf.argmax(inputs[1], 1), [-1, 1]),
                                            num_sampled=1000,
                                            num_classes=200000)

def custom_loss(y_true, y_pred):
    return K.tf.reduce_mean(y_pred)


num_classes = 200000
input = Input(shape=(300,))
target_input = Input(shape=(num_classes,))

dense = Dense(num_classes)

outputs = dense(input)
outputs = SampledSoftmax()([outputs, target_input])

model = Model([input, target_input], outputs)
model.compile(optimizer=u'adam', loss=custom_loss)
# train as desired