Restoring Trainable Variables from Saved Model

[danilojsl]

I defined a model using Tensorflow 2 with a Model Sublassing API on Python

class BiLSTMModel(tf.keras.Model):

    def __init__(self, lstm_dims):
        super().__init__()
        self.ldims = lstm_dims
        self.blockLstm = FirstBlockLSTMModule(lstm_dims)
        self.nextBlockLstm = NextBlockLSTM(lstm_dims)

    def call(self, inputs):
        # Forward pass
        block_lstm1_output = self.blockLstm(inputs)
        block_lstm2_output = self.nextBlockLstm(block_lstm1_output)
        return block_lstm2_output

I implemented a custom training loop and after each epoch, I just save this Keras model

biLSTMModel = BiLSTMModel(lstm_dims)
for epoch in range(epochs):
 # Training code
 tf.saved_model.save(biLSTMModel, export_dir)

It generates the expected structure with assets and variables folders, along with saved_model.pb and variables files

Later on, when loading the model I can restore all trainable variables values of the model just like this:

loaded_bi_lstm = tf.saved_model.load(model_path)
infer_bi_lstm = loaded_bi_lstm.signatures["serving_default"]

# LSTM weights
w_first_lstm = tf.Variable(infer_bi_lstm.trainable_variables[0])
wig_first_lstm = tf.Variable(infer_bi_lstm.trainable_variables[1])
wfg_first_lstm = tf.Variable(infer_bi_lstm.trainable_variables[2])
wog_first_lstm = tf.Variable(infer_bi_lstm.trainable_variables[3])

The model on tensorflow-java is loading without errors.

@Test
public void shouldRestoreVariablesFromSavedModel() {
    SavedModelBundle model = SavedModelBundle.load(SAVED_MODEL_DP_PATH, SavedModelBundle.DEFAULT_TAG);
    ConcreteFunction concreteFunction = model.function("serving_default");

    //Trying to find some object that stores trainables variables data
    MetaGraphDef metaGraphDef = model.metaGraphDef();
    SignatureDef sig = metaGraphDef.getSignatureDefOrThrow("serving_default");
}

But I don't find a way to restore the trainable variables. When debugging, it seems there is no object that stores these data. Could you please guide me on how to get all trainable variables values from a saved model on Java.

Thanks

[rnett]

I'm fairly sure the variable values are serialized into the exported graph, so they should be there. Did you try running it?

Actually extracting the variables is harder, your best bet is probably to either manually inspect the exported graph or use Graph.operations() (or the equivalent in Python) to look for the variable names, then fetch-ing them from the session (which will get the value).

[karllessard]

Normally the state of your tensors should be automatically restored when loading the saved model if you have the variables folder in it. You can also try to load variables at a specific checkpoint using this method from the bundle's session.

[danilojsl]

Hi @rnett and @karllessard thanks for the feedback. I rewrote the code to load the model like this:

@Test
public void shouldRestoreVariableFromSavedModel() {
  String prefix = "/BiLSTM/variables/variables";
  SavedModelBundle model = SavedModelBundle.load(SAVED_MODEL_DP_PATH, SavedModelBundle.DEFAULT_TAG);
  Graph restoredGraph = model.graph();
  try(Session restoredSession = new Session(restoredGraph)) {
      restoredSession.restore(prefix);
      String variableToRestore = "bi_lstm_model/FirstBlockLSTMModule/w_first_lstm";
      AutoCloseableList<Tensor> tensorResult = new AutoCloseableList<>(restoredSession.runner().fetch(variableToRestore).run());
    }
}

But I got this error:

java.lang.IllegalArgumentException: No tensor type has been registered for data type DT_RESOURCE at org.tensorflow.internal.types.registry.TensorTypeRegistry.find(TensorTypeRegistry.java:50) at org.tensorflow.RawTensor.fromHandle(RawTensor.java:147) at org.tensorflow.Session.run(Session.java:702) at org.tensorflow.Session.access$100(Session.java:72) at org.tensorflow.Session$Runner.runHelper(Session.java:381) at org.tensorflow.Session$Runner.run(Session.java:329) at SavedModelTest.shouldRestoreVariableFromSavedModel(SavedModelTest.java:198)

[rnett]

Well that's odd, seems like a bug. Can you share your saved model?

If you can't, can you drop a breakpoint here and tell me what output.dataType() is? That block of code is supposed to detect when you are fetching a resource variable and fetch the value instead.

You're probably on release, aren't you. Variable fetching is available in snapshots, can you try that? If you are on snapshots, see above. If you need to stay on release, add or find a read op of the variable and fetch that instead, you can look at the code linked above.

[danilojsl]

Hi @rnett, yes I'm in release 0.3.0.

Thanks for the hints. I could make it work, I'm adding here the code in case you want to check and for future reference. I also cast it to a primitive variable to make sure the values deserialized are the same as in Python.

@Test
public void shouldRestoreVariableFromSavedModel() {
  String prefix = "dp-parser.model7/BiLSTM/variables/variables";
  SavedModelBundle model = SavedModelBundle.load(SAVED_MODEL_DP_PATH, SavedModelBundle.DEFAULT_TAG);
  Graph restoredGraph = model.graph();
  try(Session restoredSession = new Session(restoredGraph)) {
      restoredSession.restore(prefix);
      String variableToRestore = "bi_lstm_model/FirstBlockLSTMModule/wfg_first_lstm/Read/ReadVariableOp";
      AutoCloseableList<Tensor> autoCloseableList = new AutoCloseableList<>(restoredSession.runner().fetch(variableToRestore).run());
      Tensor tensor = autoCloseableList.get(0);
      TFloat32 tFloat32 = (TFloat32) tensor;
      float[] tensorRawValues = TensorValues.castToPrimitiveFloat(tFloat32);
      System.out.println(Arrays.toString(tensorRawValues));
  }
}