TensorFlow - Forming Graphs

TensorFlow - Forming Graphs with `tf.function`

In TensorFlow 2.x, eager execution is enabled by default, meaning operations are executed immediately. However, TensorFlow still supports graph execution, which offers benefits like performance optimization and portability. The primary way to create graphs in TF2.x is by using the tf.function decorator. When you decorate a Python function with tf.function, TensorFlow can trace it to create a callable TensorFlow graph.

We’ll demonstrate this by simulating a simple physical system governed by a partial differential equation (PDE) – specifically, a 2D wave equation, like ripples in a pond. We will define the simulation steps within a Python function and then use tf.function to convert it into a high-performance graph operation.

Let’s assume our pond is a 2D grid of size N x N (e.g., N=500).

Step 1: Import necessary libraries.

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import time # To compare execution speeds (optional)

# Helper function to ensure eager execution is off if TF1 compatibility was enabled elsewhere
# For pure TF2, this is not strictly needed as eager is default.
# tf.compat.v1.enable_eager_execution() # Ensure eager is on for TF2 behavior

print(f"TensorFlow version: {tf.__version__}")
print(f"Eager execution enabled: {tf.executing_eagerly()}")

Step 2: Define helper functions for convolution (to compute the Laplacian).

# Function to create a convolution kernel from a 2D array
def make_kernel(a):
    a = np.asarray(a, dtype=np.float32)
    # Reshape to [filter_height, filter_width, in_channels, out_channels]
    # For depthwise_conv2d, it's [filter_height, filter_width, in_channels, channel_multiplier]
    # Here, in_channels=1, channel_multiplier=1
    return tf.constant(a.reshape(list(a.shape) + [1, 1]))

# Simplified 2D convolution. Note: tf.nn.conv2d is more general.
# tf.nn.depthwise_conv2d is used here for simplicity matching older example logic.
def simple_conv(x, k):
    # x shape: [height, width]
    # k shape: [kernel_height, kernel_width, 1, 1]
    # Add batch and channel dimensions to x: [1, height, width, 1]
    x_expanded = tf.expand_dims(tf.expand_dims(x, 0), -1)

    # Depthwise convolution: each input channel is convolved with its own set of filters.
    # Here, in_channels=1, so it's similar to a standard convolution with 1 input channel.
    y = tf.nn.depthwise_conv2d(x_expanded, k, strides=[1, 1, 1, 1], padding='SAME')

    # Remove batch and channel dimensions: [height, width]
    return y[0, :, :, 0]

# Function to compute the 2D Laplacian of an array
def laplace(x):
    # Laplacian kernel: approximates the second spatial derivative
    laplace_k = make_kernel([
        [0.5, 1.0, 0.5],
        [1.0, -6., 1.0], # Center weight is negative sum of others for some stencils
        [0.5, 1.0, 0.5]
    ])
    return simple_conv(x, laplace_k)

Step 3: Set up initial conditions and simulation parameters.

N = 100 # Reduced grid size for faster execution in tutorial

# Initial Conditions: simulate some raindrops hitting a pond
u_init = np.zeros([N, N], dtype=np.float32)
ut_init = np.zeros([N, N], dtype=np.float32) # Initial velocity

# Add some disturbances (raindrops)
for _ in range(20): # Reduced number of raindrops
    a, b = np.random.randint(0, N, 2)
    u_init[a, b] = np.random.uniform()

# Display initial state of the pond
plt.figure(figsize=(6,6))
plt.imshow(u_init)
plt.title("Initial State of the Pond (u_init)")
plt.colorbar()
plt.show()

# Simulation parameters (will be arguments to our tf.function)
# eps: time resolution (delta_t)
# damping: wave damping factor

# Create TensorFlow Variables for the simulation state
U = tf.Variable(u_init, name='U')       # Current displacement
Ut = tf.Variable(ut_init, name='Ut')    # Current velocity

Step 4: Define the PDE update rule within a Python function and decorate it with tf.function.

@tf.function
def pde_step(current_U, current_Ut, eps, damping):
    """Performs one step of the PDE simulation."""
    # Discretized PDE update rules (Euler integration for wave equation)
    # U_t+1 = U_t + eps * Ut_t
    # Ut_t+1 = Ut_t + eps * (laplacian(U_t) - damping * Ut_t)

    new_U = current_U + eps * current_Ut
    new_Ut = current_Ut + eps * (laplace(current_U) - damping * current_Ut)

    # Update the state variables (assign new values)
    current_U.assign(new_U)
    current_Ut.assign(new_Ut)
    # tf.group is not needed in TF2 eager/tf.function as assignments are sequenced
    return current_U, current_Ut

# Test one step eagerly (optional, to see if it works before many iterations)
# pde_step(U, Ut, tf.constant(0.03, dtype=tf.float32), tf.constant(0.04, dtype=tf.float32))
# print("U after one eager step (max value):", tf.reduce_max(U).numpy())

Step 5: Run the simulation.

num_iterations = 1000
eps_val = tf.constant(0.03, dtype=tf.float32)
damping_val = tf.constant(0.04, dtype=tf.float32)

print("Starting simulation using tf.function...")
start_time = time.time()

for i in range(num_iterations):
    # Call the tf.function-decorated step function
    U, Ut = pde_step(U, Ut, eps_val, damping_val)

    # Visualize progress occasionally
    if i % 200 == 0 or i == num_iterations -1: # Reduced frequency for faster overall run
        print(f"Iteration {i}")
        plt.figure(figsize=(6,6))
        plt.imshow(U.numpy()) # Use .numpy() to get array for matplotlib
        plt.title(f"Pond State at Iteration {i}")
        plt.colorbar()
        plt.show()

end_time = time.time()
print(f"Simulation finished in {end_time - start_time:.2f} seconds.")

The output will consist of plots showing the initial state of the ‘pond’ and its state at various iterations during the simulation. You should observe wave-like patterns propagating and dissipating. The tf.function decorator converts pde_step into a TensorFlow graph for potentially faster execution, especially for computations involving many small TensorFlow operations, by reducing Python overhead and enabling graph optimizations.

Key takeaways for tf.function:

It bridges the gap between eager execution’s ease of use and graph execution’s performance.
It works best with TensorFlow operations. Python-native operations or NumPy operations inside a tf.function might get re-traced frequently or lead to tf.py_function calls, which can hinder performance.
Input arguments to a tf.function that are Python scalars or NumPy arrays might trigger re-tracing if their values change. Using tf.Tensor arguments is generally preferred for stability.
tf.Variables can be updated within a tf.function using their .assign() method.

For more details on tf.function, refer to the official TensorFlow guide: https://www.tensorflow.org/guide/function

TensorFlow - Forming Graphs

TensorFlow - Forming Graphs with tf.function

TensorFlow - Forming Graphs with `tf.function`