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

时间:2017-12-19 17:35:17

标签: tensorflow deep-learning keras loss-function

我正在Keras训练语言模型,并希望通过使用采样softmax作为我网络中的最终激活功能来加速训练。从TF文档中,我看起来需要为weightsbiases提供参数,但我不确定这些参数的输入是什么。好像我可以在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参数传递吗?或者这需要写成最终完全连接层之后的层。这里的任何指导将不胜感激。感谢。

1 个答案:

答案 0 :(得分:2)

这里的主要观察结果是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