Inside TensorFlow: Tensors, Graphs, and AutoDiff – A Deep Dive

TensorFlow is a powerful open-source library developed by Google for numerical computation and large-scale machine learning. At the heart of TensorFlow are two core concepts – tensors and computational graphs.

Understanding these abstractions and how they work under the hood is key to writing efficient and scalable TensorFlow code. In this post, we‘ll dive deep into the technical details of tensors and graphs, exploring how they are represented in memory, executed on hardware, and automatically differentiated. We‘ll also look at some advanced TensorFlow libraries and techniques for working with complex graphs.

By the end of this post, you‘ll have a expert-level understanding of TensorFlow‘s core primitives and how to leverage them to build sophisticated machine learning models. Let‘s get started!

Tensors: The Anatomy of a Multi-Dimensional Array

The central unit of data in TensorFlow is the tensor – a multi-dimensional array of numerical values. Tensors are a generalization of vectors and matrices to higher dimensions.

Mathematically, a tensor is defined as a set of primitive values (e.g. integers, floats, booleans) shaped into an array with any number of dimensions. A tensor has the following properties:

• Rank: The number of dimensions of the tensor (e.g. 0, 1, 2, 3…)
• Shape: The size of each dimension of the tensor
• Data type: The primitive data type contained in the tensor (e.g. int32, float32, bool)

Here are some examples of tensors of different ranks:

# Rank 0 (scalar)  
x = 3
# x is a tensor of shape [] and rank 0

v = [1, 2, 3] 


m = [[1, 2, 3],
[4, 5, 6]] 


t = [[[1], [2], [3]],
[[4], [5], [6]]]

Tensors can theoretically have any rank, although in practice ranks 0 through 4 are most common. In TensorFlow, tensors are implemented as instances of the tf.Tensor class.

Tensor Data Types

Each tensor has a single primitive data type that is set when the tensor is created. Some common tensor types are:

• tf.float32: 32-bit floating point
• tf.int32: 32-bit signed integer
• tf.uint8: 8-bit unsigned integer
• tf.bool: Boolean

The data type determines the byte size of each element and the concrete interpretation of the bits. You can check the type of a tensor using the dtype attribute:

x = tf.constant([1.5, 2.5, 3.5], dtype=tf.float32)
print(x.dtype)  # prints "tf.float32"

It‘s important to be mindful of data types to ensure your tensors have compatible shapes and to avoid under/overflows and numerical precision issues.

Tensor Storage and Striding

Under the hood, a tensor stores its primitive values in a contiguous block of memory. The memory layout is determined by the shape and strides of the tensor.

The strides of a tensor define the number of bytes to skip in each dimension when traversing the values. This allows tensors to use memory efficiently by avoiding unnecessary copying when slicing or transposing dimensions.

Lets look at an example to illustrate how striding works:

  
t = [[1 2 3]
     [4 5 6]]  

In general, the byte offset of a tensor element is calculated as:

offset + index[0] * stride[0] + index[1] * stride[1] + ... + index[n-1] * stride[n-1] 

Where index is a vector of 0-based indices into each dimension, and stride is a vector of byte strides for each dimension.

TensorFlow tensors store their data in row-major order and automatically calculate strides based on the shape. You can access the strides of a tensor using the tf.Tensor.get_shape() method.

Sparse Tensors

In some cases, tensors may have many empty or default values. Storing these tensors densely is memory inefficient. For this reason, TensorFlow provides the tf.SparseTensor class to efficiently store sparse tensors.

A SparseTensor only stores the non-empty values and their indices. Internally, it is comprised of three dense tensors:

• indices: A 2-D int64 tensor of shape [N, ndims] specifying the indices of the non-empty values, where N is the number of values and ndims is the rank of the sparse tensor.
• values: A 1-D tensor containing the non-empty values corresponding to indices
• dense_shape: A 1-D int64 tensor specifying the shape of the dense tensor

Here‘s how you would create a sparse tensor:

sp = tf.SparseTensor(indices=[[0, 0], [1, 2]],
                     values=[1, 2],
                     dense_shape=[3, 4])

This creates a rank-2 sparse tensor with two non-empty values, 1 at index [0, 0] and 2 at index [1, 2], and a dense shape of [3, 4].

Sparse tensors can be used in many of the same operations as dense tensors, such as addition, multiplication, and slicing. TensorFlow will automatically convert between sparse and dense tensors as needed.

Computational Graphs: Encoding a Neural Network

Now that we have a solid grasp on tensors, let‘s look at how they are used in TensorFlow programs. The core idea of TensorFlow is to define a dataflow graph – a directed acyclic graph (DAG) specifying a series of operations to perform on tensors.

A TensorFlow graph consists of two main components:

• Nodes represent operations (ops) to perform on tensors, such as addition, multiplication, convolution, etc. Each node takes zero or more tensors as inputs and produces zero or more tensors as output.

• Edges represent the tensors that flow between nodes. An edge connects one node‘s output to another node‘s input.

Let‘s look at a toy example of defining a simple neural network graph:

import tensorflow as tf

x = tf.placeholder(tf.float32, shape=[None, 784], name="x")
y = tf.placeholder(tf.float32, shape=[None, 10], name="y")

W = tf.Variable(tf.zeros([784, 10]), name="W") 
b = tf.Variable(tf.zeros([10]), name="b")

z = tf.matmul(x, W) + b
probabilities = tf.nn.softmax(z)

loss = tf.reduce_mean(
tf.nn.softmax_cross_entropy_with_logits(labels=y, logits=z))

optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.5)
train_op = optimizer.minimize(loss)

This code snippet defines a graph for a simple feedforward neural network with a single hidden layer. Let‘s break it down:

• First we define placeholder tensors for feeding in input features x and labels y when we run the graph. Placeholders are special tensors whose values must be supplied at runtime.

• Next we define Variable tensors for the weights W and biases b of the network. Variables are stateful tensors that can be updated during training.

• We then define the forward pass of the network, which consists of a matrix multiplication of the inputs and weights (tf.matmul()) followed by a bias addition and softmax activation (tf.nn.softmax()). This produces a tensor of predicted class probabilities.

• We define a cross-entropy loss function (tf.nn.softmax_cross_entropy_with_logits()) to measure the discrepancy between the predicted and true labels, and use tf.reduce_mean() to compute the mean loss over the batch.

• Finally, we define an optimizer (tf.train.GradientDescentOptimizer) and a training op (optimizer.minimize(loss)) to adjust the weights and biases to minimize the loss.

The key thing to notice is that none of these ops actually perform any computation. They simply define a symbolic graph structure that specifies what computation should occur when the graph is run. We can visualize this graph in TensorBoard:

Each node in the graph represent an operation, and the edges represent the flow of tensor data between ops. The blue nodes are the placeholders where we feed in data, and the orange nodes are the variables.

Notice there are also many more nodes than the ones we defined explicitly, such as MatMul, Add, Sum, Mul, Sub, etc. These are lower-level ops that are automatically created by TensorFlow to implement the higher-level ops we defined, such as tf.matmul, tf.nn.softmax_cross_entropy_with_logits, tf.reduce_mean, etc.

Running a TensorFlow Graph

To actually execute the neural network training, we need to run the graph in a TensorFlow Session:

with tf.Session() as sess:
  # Initialize variables  
  sess.run(tf.global_variables_initializer())

for _ in range(1000):
batch_x, batch_y = next_batch(100) 
sess.run(train_op, feed_dict={x: batch_x, y: batch_y})

correct_prediction = tf.equal(tf.argmax(probabilities, 1), tf.argmax(y, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
print("Test accuracy:", sess.run(accuracy, feed_dict={x: test_x, y: test_y}))

Here‘s what‘s happening:

1. We create a Session object within a with block. The session acts as a runtime environment for executing a graph.

2. We initialize the Variable tensors in our graph by running tf.global_variables_initializer(). This assigns initial values to the weights and biases.

3. We run a training loop for 1000 steps. Each step, we fetch the next minibatch of data, then run the train_op to perform a gradient descent update on the weights and biases, feeding the minibatch into the x and y placeholders.

4. After training, we evaluate the accuracy of the model on a test set. We run the accuracy tensor, which returns the fraction of correct predictions.

One of the advantages of TensorFlow graphs is that the same code can be run without modification on CPUs, GPUs, or even distributed across multiple machines. The TensorFlow runtime takes care of allocating the operations to available devices and parallelizing the computation.

Automatic Differentiation on Graphs

Another key benefit of dataflow graphs is they enable easy automatic differentiation. In order to train a neural network, we need to compute the gradients of the loss function with respect to all the parameters (e.g. W and b).

With TensorFlow, we can compute these gradients automatically using the tf.gradients() function:

W_grad, b_grad = tf.gradients(loss, [W, b])

This create nodes in the graph to compute the partial derivative of loss with respect to W and b using the chain rule of differentiation. Under the hood, TensorFlow performs a technique called reverse-mode autodiff or backpropagation on the graph.

The key idea is to first perform a forward pass through the graph, computing the output of each node and caching intermediate values. Then, we perform a backward pass, computing the "adjoint" of each node in reverse order. The adjoint measures how much each node contributes to the final output. By applying the chain rule recursively, we can compute the gradient with respect to each parameter.

TensorFlow‘s autodiff capabilitites are very general and flexible. You can take gradients of any scalar tensor with respect to any other tensor, including taking higher-order derivatives. This allows for a wide range of optimization and inference algorthms to be easily implemented.

Conclusion

In this post, we took a deep dive into the guts of TensorFlow, exploring how it represents data as tensors and computation as dataflow graphs. We saw how:

• Tensors are multi-dimensional arrays that efficiently store data in contiguous memory blocks.
• Sparse tensors can be used to store tensors with many empty values
• Computation graphs define a series of operations to perform on tensors.
• Graphs enable easy parallelization and distribution of computation
• Autodiff on graphs allows for easy computation of gradients.

I encourage you to explore these concepts further by building your own neural networks in TensorFlow. The TensorFlow tutorials are a great resource to get started.

Some other topics worth investigating are:

• TensorFlow Fold for dynamic batching of computation graphs
• TensorFlow XLA for compiling graphs for specialized hardware
• TensorFlow Probability for probabilistic programming with graphs
• TensorNetwork for working with tensor networks in quantum computing.

I hope this post gave you a solid understanding of TensorFlow‘s core abstractions and internals. While TensorFlow‘s API can seem complex at first, its core concepts of tensors and computation graphs are actually quite simple and elegant. Mastering them will give you the tools you need to build state-of-the-art machine learning models. Happy coding!