TensorFlow: Save and Load to Continue Training (tf.Module instead of tf.keras.Model)

After saving a tf.Module with tf.saved_model.save, the model can be loaded by tf.saved_model.load, the model can be train more by applying gradients to:

• M.Some_Layer.trainable_variables
• M.Some_Var

The list Model.trainable_variables is no longer in the model after loading.

Source code:

%tensorflow_version 2.x
%reset -f

#libs
import tensorflow as tf;
from tensorflow.keras.layers import *;

#constants
BSIZE = 4;

#model
class model(tf.Module):
  def __init__(this):
    super().__init__();
    #this.W1 = tf.Variable(tf.random.uniform([2,20], -1,1));
    #this.B1 = tf.Variable(tf.random.uniform([  20], -1,1));
    this.Layer1 = Dense(20, activation=tf.nn.leaky_relu);

    this.W2 = tf.Variable(tf.random.uniform([20,1], -1,1));
    this.B2 = tf.Variable(tf.random.uniform([   1], -1,1));

  @tf.function(input_signature=[tf.TensorSpec([BSIZE,2], tf.float32)])
  def __call__(this,Inp):
    #H1  = tf.nn.leaky_relu(tf.matmul(Inp,this.W1) + this.B1);
    H1  = this.Layer1(Inp);
    Out = tf.sigmoid(tf.matmul(H1,this.W2) + this.B2);    
    return Out;

#data (OR)
X = tf.convert_to_tensor([[0,0],[0,1],[1,0],[1,1]], tf.float32);
Y = tf.convert_to_tensor([[0],  [1],  [1],  [1]  ], tf.float32);

#train
Model = model();
Loss  = tf.losses.MeanSquaredError();
Optim = tf.optimizers.SGD(1e-1);
Steps = 10;

for I in range(Steps):
  if I%(Steps/10)==0:
    Out = Model(X);
    Lv  = Loss(Y,Out);
    print("Loss:",Lv.numpy());

  with tf.GradientTape() as T:
    Out = Model(X);
    Lv  = Loss(Y,Out);

  Grads = T.gradient(Lv, Model.trainable_variables);
  Optim.apply_gradients(zip(Grads, Model.trainable_variables));

Out = Model(X);
Lv  = Loss(Y,Out);
print("Loss:",Lv.numpy(),"(Last)");

#save
print("\nSaving model...");
Dir = "/tmp/models/test";
tf.saved_model.save(Model,Dir);

#load
print("\nLoading model...");
M = tf.saved_model.load(Dir);
print(vars(M).keys());
print(tf.keras.backend.flatten(M(X)).numpy());

#train more
print("\nContinue training...");
Steps = 1000;

for I in range(Steps):
  if I%(Steps/10)==0:
    Out = M(X);
    Lv  = Loss(Y,Out);
    print("Loss:",Lv.numpy());

  with tf.GradientTape() as T:
    Out = M(X);
    Lv  = Loss(Y,Out);

  Grads = T.gradient(Lv, M.Layer1.trainable_variables+[M.W2,M.B2]);
  Optim.apply_gradients(zip(Grads, M.Layer1.trainable_variables+[M.W2,M.B2]));

Out = M(X);
Lv  = Loss(Y,Out);
print("Loss:",Lv.numpy(),"(Last)");
print(tf.keras.backend.flatten(M(X)).numpy());
#eof