Can Keras models be trained using the TensorFlow C API?

Viewed 846

I have a model defined in Keras that I would like to train using the TensorFlow C API for a reinforcement learning application. (Actually by way of the Rust programming language, but it uses the C API directly.) I have struggled to find some way to do this, but as of yet no luck. I had expected ops or functions to be serialized with a SavedModel which would allow training, but I don't see them. Is it even possible to train Keras models from C? Or TF models in general?

1 Answers

You can create a Custom Keras Model with one tf.function for training and another for predictions.

class MyModel(tf.keras.Model):
def __init__(self, name=None, **kwargs):
    super().__init__(**kwargs)
    
    #### define your layers here ####
    self.vec_layer = vectorizer
    self.preds = model

#### override call function (used for prediction) ####
@tf.function
def call(self, inputs, training=False): 
    #### define model structure ####
    x = self.vec_layer(inputs)
    return self.preds(x)
    #return {'output': self.preds(x)} #return labeled result

#### this function only calls the train_step. ####
#It can return whatever the user wants. Examples returning wither loss or nothing
@tf.function
def training(self, data):
    loss = self.train_step(data)['loss']
    return {'loss': loss}
    return {}

model = MyModel(name="end_model")

Then, you have to compile and build your model.

model.compile(loss="sparse_categorical_crossentropy", optimizer="SGD", metrics=["acc"])
model.fit([["Ok."]], [1], batch_size=1, epochs=1)

Before saving the model, you need to define your tf.functions as concrete functions and pass them as signature functions.

call_output = model.call.get_concrete_function(tf.TensorSpec([None,1], tf.string, name='input'))
train_output = model.training.get_concrete_function((tf.TensorSpec([None,1], tf.string, name='inputs'),tf.TensorSpec([None,1], tf.float32, name='target')))
model.save("model_saved.tf", save_format="tf", signatures={'train': train_output, 'predict': call_output})

Now, you can train/predict in your model using the C API

TF_Graph* graph = TF_NewGraph();
TF_Status* status = TF_NewStatus();
TF_SessionOptions* SessionOpts = TF_NewSessionOptions();
TF_Buffer* RunOpts = NULL;

const char* saved_model_dir = "MODEL_DIRECTORY";
const char* tags = "serve"; //saved_model_cli

int ntags = 1;
TF_Session* session = TF_LoadSessionFromSavedModel(SessionOpts, RunOpts, saved_model_dir, &tags, ntags, graph, NULL, status);

if(TF_GetCode(status) == TF_OK)
{
    printf("TF_LoadSessionFromSavedModel OK\n");
}
else
{
    printf("%s",TF_Message(status));
}

Then, create your input and output Tensors as here. Now, just call TF_SessionRun.

#### predicting ####
TF_SessionRun(session, nullptr,
              input_operator, input_values, 1,
              output_operator, output_values, 1,
              nullptr, 0, nullptr, status);

#### training ####
TF_SessionRun(session, nullptr, input_target_operators, input_target_values, 2, loss_operator, loss_values, 1, nullptr, 0, nullptr, status);

Important to note:

  • In order to get the variable names (input, target, output, etc) you can use the "saved_model_cli" command, explained here.
  • In theory, you would have to pass the signature_def name of the function that you want to execute with the "TF_SessionRun" method. However, I could not make this work. However, somehow, passing the right inputs and output values for the "TF_SessionRun" makes it automatically select the right method to use.
  • The code above is part of my project so it won't compile on its own. However, I hope it will assist you, as I am not with enough time to make a whole functional example.
  • Other than the links already mentioned, this explains how to train the model using TF v1, and this is a more complete example of how you can construct and save your model in python.
Related