A Full-feature Low-level TensorFlow Programme (Update)

The following code implements a full-feature low-level TensorFlow programme with:

• Model based on tf.Module
• A sample Keras layer tf.keras.layers.Dense
• A sample raw TensorFlow layer with tf.Variable
• Initialisation of layers with tf.random.uniform
• Use of rectifier and sigmoid-like in model
• Autograph with @tf.function
• Load and normalise data
• Create a batch generator (tf.data.Dataset)
• Sample user of loss function and optimiser
• Train model (adjust weights) loop with gradient calculation and applying gradient
• A simple way to calculate average loss after a number of batches
• Save model to a directory with tf.saved_model.save
• Load a model from a directory with tf.saved_model.load
• Continue training the model
• Infer the training data with the trained model.

Source code:

%tensorflow_version 2.x

%reset -f

#libs

import tensorflow as tf;

from tensorflow.keras.layers import *;

#constants

DSIZE = 4;

BSIZE = 2;

#model

class model(tf.Module):

  #constructor

  def __init__(this):

    #utilise a keras layer

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

    #raw tf layer

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

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

    #single neuron output layer, also raw

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

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

  #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.sigmoid(tf.matmul(H2,this.W3) + this.B3); 

    return Out;

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

#data (XOR data for example)

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

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

#normalise data to range [0,1]

X = X/tf.reduce_max(X);

Y = Y/tf.reduce_max(Y);

#make batch generator

Data = tf.data.Dataset.from_tensor_slices((X,Y));

Data = Data.repeat().shuffle(DSIZE).batch(BSIZE,drop_remainder=True);

#train

Model = model();

Loss  = tf.losses.MeanSquaredError();

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

Steps = 1000;

Laft  = 100;

Lsum  = 0;

Iter  = iter(Data);

for I in range(Steps):

  Batch = next(Iter);

  Inp   = Batch[0]; #input

  Exp   = Batch[1]; #expected

  with tf.GradientTape() as T:

    Lval  = Loss(Exp,Model(Inp));

    Lsum += Lval.numpy();

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

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

  if I%Laft==Laft-1:

    print("Average Loss:",Lsum/Laft);

    Lsum = 0;

#save

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

Model_Dir = "/tmp/my-model";

tf.saved_model.save(Model,Model_Dir);

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

#load back model

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

Model = tf.saved_model.load(Model_Dir);

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

#continue training the model

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

for I in range(Steps):

  Batch = next(Iter); #load another dataset or use back the infinite Iter above

  Inp   = Batch[0];

  Exp   = Batch[1];

  with tf.GradientTape() as T:

    Lval  = Loss(Exp,Model(Inp));

    Lsum += Lval.numpy();

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

  Grads  = T.gradient(Lval, Params);

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

  if I%Laft==Laft-1:

    print("Average Loss:",Lsum/Laft);

#infer the training data

print("\nInferring the training data...");

Data = tf.data.Dataset.from_tensor_slices((X,Y));

Data = Data.batch(BSIZE,drop_remainder=True);

for Inp,Exp in Data:

  print("Input:");

  print(Inp.numpy());

  print("Expected:");

  print(Exp.numpy());

  print("Prediction:");

  print(Model(Inp).numpy());

  print();

#eof

Result:
Average Loss: 0.10872401443310081 Average Loss: 0.01269742899108678 Average Loss: 0.00553579420549795 Average Loss: 0.0033595020952634515 Average Loss: 0.0023345320229418575 Average Loss: 0.001752474318491295 Average Loss: 0.0013855347724165768 Average Loss: 0.0011406122986227273 Average Loss: 0.0009637939184904098 Average Loss: 0.0008266885939519853 Saving model... INFO:tensorflow:Assets written to: /tmp/my-model/assets total 36 drwxr-xr-x 2 root root 4096 Dec 26 10:57 assets/ drwxr-xr-x 2 root root 4096 Dec 26 11:06 variables/ drwxr-xr-x 4 root root 4096 Dec 26 11:06 ./ drwxrwxrwt 1 root root 4096 Dec 26 11:06 ../ -rw-r--r-- 1 root root 18839 Dec 26 11:06 saved_model.pb Loading back model... dict_keys(['_self_setattr_tracking', '_self_unconditional_checkpoint_dependencies',
'_self_unconditional_dependency_names', '_self_unconditional_deferred_dependencies',
'_self_update_uid', '_self_name_based_restores',
'Layer1', 'W2', 'B2', 'W3', 'B3',
'signatures', '__call__', 'tensorflow_version', 'tensorflow_git_version']) Continue training... Average Loss: 0.000725053078494966 Average Loss: 0.001368767161620781 Average Loss: 0.0019466301079955884 Average Loss: 0.0024687587827793324 Average Loss: 0.002944960643653758 Average Loss: 0.00338194992946228 Average Loss: 0.003784944185754284 Average Loss: 0.0041586664071655835 Average Loss: 0.004504695165960584 Average Loss: 0.0048269858321873475 Inferring the training data... Input: [[0. 0.] [0. 1.]] Expected: [[0.] [1.]] Prediction: [[0.01820062] [0.98395973]] Input: [[1. 0.] [1. 1.]] Expected: [[1.] [0.]] Prediction: [[0.9820526] [0.0182383]]