A Full-feature Low-level TensorFlow Model

A low-level TensorFlow model should be created with tf.Module, and for the ease of save and load, use tf.saved_model APIs.

Beside the raw tf.Variable(s), a full-feature TensorFlow model should contain Keras layers too: Dense, ConvXD, LSTM, MaxPoolXD, RNN, GRU, etc.

Remember to use the decorator @tf.function to convert Python code into graph in C++ for performance.

Data feeding is also a matter in model training performance, use APIs in  tf.data module for this.

Source code:

%tensorflow_version 2.x

%reset -f

#libs

import tensorflow as tf;

from tensorflow.keras.layers import *;

#global constants

BSIZE=4;

#model

class model(tf.Module):

  #constructor

  def __init__(this):

    #example Keras layer to use inside tf.Module

    this.Layer1 = Dense(30, tf.nn.leaky_relu); 

    #raw params

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

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

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

    this.B3 = tf.Variable(tf.random.uniform([    1], -1,1, tf.float32));

  #model call

  @tf.function(input_signature=[tf.TensorSpec([BSIZE,2], tf.float32)])

  def __call__(this,Inp):

    H1  = this.Layer1(Inp);

    H2  = tf.nn.leaky_relu(tf.matmul(H1,this.W2) + this.B2);

    Out = tf.nn.leaky_relu(tf.matmul(H2,this.W3) + this.B3);

    return Out;

#PROGRAMME ENTRY POINT==========================================================

#data

X = tf.constant([[0,0],[0,1],[1,0],[1,1]], tf.float32);

Y = tf.constant([[0  ],[1  ],[1  ],[0  ]], tf.float32);

#train

Model = model();

Loss  = tf.losses.MeanSquaredError();

Optim = tf.optimizers.SGD(1e-1);

Steps = 20;

print("Training model...");

for I in range(Steps):

  if I%(Steps/10)==0:

    Lv = Loss(Y,Model(X));

    print("Loss:",Lv.numpy());

  with tf.GradientTape() as T:

    Lv = Loss(Y,Model(X));

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

  Optim.apply_gradients(zip(Grads, Model.trainable_variables));

#get final loss

Lv = Loss(Y,Model(X));

print("Loss:",Lv.numpy(),"(Final)");

#save model

print("\nSaving model...");

Dir = "/tmp/my-model";

tf.saved_model.save(Model,Dir);

#note that the saved_model.pb is not frozen (containing constants)

#its param values are still in /tmp/my-model/variables

%ls -laX --group-directories-first /tmp/my-model

#load back model

print("\nLoading back model...");

Model = tf.saved_model.load(Dir);

print(vars(Model).keys());

#continue training

print("\nContinue training...");

for I in range(Steps):

  if I%(Steps/10)==0:

    Lv = Loss(Y,Model(X));

    print("Loss:",Lv.numpy());

  with tf.GradientTape() as T:

    Lv = Loss(Y,Model(X));

  #this is a little bit different from the loop above

  #only Keras layers now have trainable_variables, list out the rest manually

  Params = Model.Layer1.trainable_variables+[Model.W2,Model.B2,Model.W3,Model.B3];

  Grads = T.gradient(Lv, Params);

  Optim.apply_gradients(zip(Grads, Params));

#get final loss

Lv = Loss(Y,Model(X));

print("Loss:",Lv.numpy(),"(Final)");

#eof