title: How can I use TensorFlow's sampled softmax loss function in a Keras model?
body: <p>I'm training a language model in Keras and would like to speed up training by using sampled softmax as the final activation function in my network. From the TF docs, it looks like I need to supply arguments for <code>weights</code> and <code>biases</code>, but I'm unsure of what is expected as input for these. It seems like I could write a custom function in Keras as follows:</p>
  
  <pre><code>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)
  </code></pre>
  
  <p>However, I'm unsure of how to "plug this in" to my existing network. The architecture for the LM is pretty dead-simple:</p>
  
  <pre><code>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')
  </code></pre>
  
  <p>Given this architecture, could I pass the <code>sampled_softmax</code> function as the <code>loss</code> argument when calling the compile method on the model? Or do this need to be written as a layer that comes after the final fully-connected layer. Any guidance here would be greatly appreciated. Thanks.</p>