Can a Tensorflow variable be trained using the Tensorflow Keras functional API model? Can a Tensorflow operation be used in the functional API Model?

Question

I am wondering if Keras model compile/training with the functional API train variables defined by tf.get_variable? Can Keras training also incorporate Tensorflow operations?

So basically I am looking to define a Keras model with Tensorflow variables and operations, then use

model = tf.keras.Model(inputs=inputs, outputs=predictions)
model.compile(optimizer=optimizer, loss=loss)
model.fit(data, labels, batch_size=batch_size, epochs=epochs)

To train the model. The reason for this is that Google's TPUs require either a Keras or TF.Estimator API, with Keras being more recommended, so I am looking to see how easily I can convert my model.

BackGround

It looks like since Tensorflow is the backend, there are ways to mix Keras/Tensorflow variables. This blog post shows how Keras variables are trained using a Tensorflow graph/session https://blog.keras.io/keras-as-a-simplified-interface-to-tensorflow-tutorial.html

from keras.layers import Dropout
from keras import backend as K

img = tf.placeholder(tf.float32, shape=(None, 784))
labels = tf.placeholder(tf.float32, shape=(None, 10))

x = Dense(128, activation='relu')(img)
x = Dropout(0.5)(x)
x = Dense(128, activation='relu')(x)
x = Dropout(0.5)(x)
preds = Dense(10, activation='softmax')(x)

loss = tf.reduce_mean(categorical_crossentropy(labels, preds))

train_step = tf.train.GradientDescentOptimizer(0.5).minimize(loss)
with sess.as_default():
    for i in range(100):
        batch = mnist_data.train.next_batch(50)
        train_step.run(feed_dict={img: batch[0],
                                  labels: batch[1],
                                  K.learning_phase(): 1})

acc_value = accuracy(labels, preds)
with sess.as_default():
    print acc_value.eval(feed_dict={img: mnist_data.test.images,
                                    labels: mnist_data.test.labels,
                                    K.learning_phase(): 0})

And also here it shows that Tensorflow variables can be used as input to a Keras model

How to set the input of a Keras layer of a functional model, with a Tensorflow tensor?

tf_embedding_input = ...    # pre-processing output tensor

# Keras model
model = Sequential()
model.add(Input(tensor=tf_embedding_input)) 
model.add(Embedding(max_features, 128, input_length=maxlen))

So I am wondering if Keras can train Tensorflow variables.

Example

I would like to train the embedding and softmax variables in the Tensorflow architecture below

  embeddings = tf.get_variable( 'embeddings', 
    initializer= tf.random_uniform([vocabulary_size, embedding_size], -1.0, 1.0))

  softmax_weights = tf.get_variable( 'softmax_weights',
    initializer= tf.truncated_normal([vocabulary_size, embedding_size],
                         stddev=1.0 / math.sqrt(embedding_size)))
  
  softmax_biases = tf.get_variable('softmax_biases', 
    initializer= tf.zeros([vocabulary_size]),  trainable=False )

  embed = tf.nn.embedding_lookup(embeddings, train_dataset) #train data set is

  embed_reshaped = tf.reshape( embed, [batch_size*num_inputs, embedding_size] )
  
  segments= np.arange(batch_size).repeat(num_inputs)

  averaged_embeds = tf.segment_mean(embed_reshaped, segments, name=None)

  loss = tf.reduce_mean(
    tf.nn.sampled_softmax_loss(weights=softmax_weights, biases=softmax_biases, inputs=averaged_embeds,
                               labels=train_labels, num_sampled=num_sampled, num_classes=vocabulary_size))

Since Tensorflow Keras uses a Tensorflow backend, I'm guessing it's somehow possible to use and train Tensorflow variables and use Tensorflow operations in training.

Why do I want to do this?

Google's TPUs require that your architecture be implemented via the Estimator API or Keras API. Since the Keras API is more recommended, there is probably interest in converting a regular Tensorflow Graph/Session to use the Keras API with as few alterations to their code as possible.

Knowing how to incorporate Tensorflow operations and train Tensorflow variables using the Keras model compile/train would greatly help with this.

Answer 1

_{Little background:}

As we know Keras is a model-level library, providing high-level building blocks for developing deep learning models.

The most important thing: Keras API does not handle tensor operations. It needs a well-optimized tensor manipulation library to do so, know as a "backend engine" for Keras.

At this time, Keras has three backend engines available: the TensorFlow backend (Google), the Theano backend, and the CNTK backend (MSFT).

Knowing how to incorporate Tensorflow operations and train Tensorflow variables using the Keras model compile/train would greatly help with this.

The only thing you should ask yourself, is what is the difference between the Keras variable and regular Tensorflow variable.

Happens to be that Keras variable have metadata. So in order to use the TensorFlow variables in Keras you convert them.

Note: A TensorFlow variable scope will have no effect on a Keras layer or model.

Finally variable sharing can be done by initializing the Keras layer (or model).

Answer 2

Would this solution help?

keras add external trainable variable to graph

You could feed your embeddings and softmax layers into the Keras model using

model.add()

and then define those variables as trainable using

model.layers[-1].trainable_weights.extend()