A cheat sheet for custom TensorFlow layers and models

This article is about compiling all those TensorFlow manuals into one page. Image from Unsplash by Patrick Tomasso.
TensorFlow v2.5.0.

Numpy mappings to TensorFlow

  • Keep in mind for all operations: in TensorFlow, the first axis is the batch_size.
  • Trick to get past the batch size when it’s inconvenient: wrap linear algebra in tf.map_fn , e.g. converting vectors of size n in a batch (batch_size x n) to diagonal matrices of size (batch_size x n x n) :
mats = tf.map_fn(lambda vec: tf.linalg.tensor_diag(vec), vecs)
  • Sum an array: np.sum(x)tf.map.reduce_sum(x)
  • Dot product of two vectors: np.dot(x,y)tf.tensordot(x,y,1)
  • Matrix product with vector A.b : np.dot(A,b)tf.linalg.matvec(A,b)
  • Trigonometry: np.sintf.math.sin
  • Transpose (batch_size x n x m)(batch_size x m x n) : np.transpose(x, axes=(0,2,1))tf.transpose(x, perm=[0,2,1])
  • Vector to diagonal matrix: np.diag(x)tf.linalg.tensor_diag(x)
  • Concatenate matrices: np.concatenate((A,B),axes=0)tf.concat([A,B],1)
  • Matrix flatten: [[A,B],[C,D]] into a single matrix:
tmp1 = tf.concat([A,B],2)
tmp2 = tf.concat([C,D],2)
mat = tf.concat([tmp1,tmp2],1)
  • Kronecker product of a vector vec of size n (makes a matrix n x n): tf.tensordot(vec,vec,axes=0) except the vectors usually come in a batch vecs of size (batch_size x n) , so we need to map:
mats = tf.map_fn(lambda vec: tf.tensordot(vec,vec,axes=0),vecs)
  • Zeros: np.zeros(n)tf.zeros_like(n) . Note that to avoid the error “Cannot convert a partially known TensorShape to a Tensor”, you should use tf.zeros_like instead of tf.zeros since the former does not require the size to be known until runtime, see also here.

Part assignment by indexes

a = np.random.rand(3)
a[0] = 5.0
a = tf.Variable(initial_value=np.random.rand(3),dtype='float32')
a[0] = 5.0
# TypeError: 'ResourceVariable' object does not support item assignment
e = tf.one_hot(indices=0,depth=3,dtype='float32')# <tf.Tensor: shape=(3,), dtype=float32, numpy=array([1., 0., 0.], dtype=float32)>
a = a + (new_val - old_val) * e

Part assignment by indexes for matrices, etc.

Getting the shape of things

What to subclass — tf.keras.layers.Layer or tf.keras.Model?

What do I need to implement to subclass a Layer?

  • Obviously call super in the constructor.
  • Obviously implement the call method.
  • You should implement get_config and from_config — they are needed often for saving the layer, e.g. if save_traces=False in the model.
  • @tf.keras.utils.register_keras_serializable(package="my_package") should be added at the top of the class. From the docs:
loaded_1 = keras.models.load_model(
"my_model", custom_objects={"MyLayer": MyLayer}

What do I need to implement to subclass a Model?

Saving and loading models and layers

model.save("model", save_traces=False)# Or just save the weights
model = tf.keras.models.load_model("model")# Or just load the weights for an already existing model

How to reduce the size of TensorBoard callbacks if you have a large custom model

logdir = os.path.join("logs",datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))tensorboard_callback = tf.keras.callbacks.TensorBoard(logdir,
val_checkpoint = tf.keras.callbacks.ModelCheckpoint(





Coding, ML, AI — oliver-ernst.com

Love podcasts or audiobooks? Learn on the go with our new app.

Recommended from Medium

How to split any string into a list using Python?

Testing Strategy in Android — Part 1

Connect to a Pulsar Cluster and Analyze Streaming NFT Data with Memgraph

Dynamics 365-integrate without requiring a licensed user

Add data to your database with a REST API in just few lines of code

C# preprocessor directives for conditional compiling

Short glimpse of the new Dabble Lab Developer Console !

Installing Anaconda for multiple users

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store
Oliver K. Ernst, Ph.D.

Oliver K. Ernst, Ph.D.

Coding, ML, AI — oliver-ernst.com

More from Medium

From scratch to CUDA installation and TensorFlow compilation from the sources on Ubuntu 20.04

Machine Learning — Researchers from the field

Fundamentals of Particle Swarm Optimization

Newton-type Methods for Optimization