Why TensorFlow NumPy Might Be the Future of Differentiable Programming

Content Overview TF NumPy and TensorFlow tf.Tensor and ND array TensorFlow interoperability Gradients and Jacobians: tf.GradientTape Trace compilation: tf.function Vectorization: tf.vectorized_map Device placement Performance comparisons Further reading TF NumPy and TensorFlow TF NumPy and TensorFlow tf.Tensor and ND array tf.Tensor and ND array TensorFlow interoperability TensorFlow interoperability Gradients and Jacobians: tf.GradientTape Gradients and Jacobians: tf.GradientTape Trace compilation: tf.function Trace compilation: tf.function Vectorization: tf.vectorized_map Vectorization: tf.vectorized_map Device placement Device placement Performance comparisons Performance comparisons Further reading Further reading TF NumPy and TensorFlow TensorFlow NumPy is built on top of TensorFlow and hence interoperates seamlessly with TensorFlow. tf.Tensor and ND array tf.Tensor ND array is an alias to tf.Tensor, so obviously they can be intermixed without triggering actual data copies. tf.Tensor x = tf.constant([1, 2]) print(x) # `asarray` and `convert_to_tensor` here are no-ops. tnp_x = tnp.asarray(x) print(tnp_x) print(tf.convert_to_tensor(tnp_x)) # Note that tf.Tensor.numpy() will continue to return `np.ndarray`. print(x.numpy(), x.numpy().__class__) x = tf.constant([1, 2]) print(x) # `asarray` and `convert_to_tensor` here are no-ops. tnp_x = tnp.asarray(x) print(tnp_x) print(tf.convert_to_tensor(tnp_x)) # Note that tf.Tensor.numpy() will continue to return `np.ndarray`. print(x.numpy(), x.numpy().__class__) tf.Tensor([1 2], shape=(2,), dtype=int32) tf.Tensor([1 2], shape=(2,), dtype=int32) tf.Tensor([1 2], shape=(2,), dtype=int32) [1 2] tf.Tensor([1 2], shape=(2,), dtype=int32) tf.Tensor([1 2], shape=(2,), dtype=int32) tf.Tensor([1 2], shape=(2,), dtype=int32) [1 2] TensorFlow interoperability An ND array can be passed to TensorFlow APIs, since ND array is just an alias to tf.Tensor. As mentioned earlier, such interoperation does not do data copies, even for data placed on accelerators or remote devices. tf.Tensor Conversely, tf.Tensor objects can be passed to tf.experimental.numpy APIs, without performing data copies. tf.Tensor tf.experimental.numpy # ND array passed into TensorFlow function. tf_sum = tf.reduce_sum(tnp.ones([2, 3], tnp.float32)) print("Output = %s" % tf_sum) # `tf.Tensor` passed into TensorFlow NumPy function. tnp_sum = tnp.sum(tf.ones([2, 3])) print("Output = %s" % tnp_sum) # ND array passed into TensorFlow function. tf_sum = tf.reduce_sum(tnp.ones([2, 3], tnp.float32)) print("Output = %s" % tf_sum) # `tf.Tensor` passed into TensorFlow NumPy function. tnp_sum = tnp.sum(tf.ones([2, 3])) print("Output = %s" % tnp_sum) Output = tf.Tensor(6.0, shape=(), dtype=float32) Output = tf.Tensor(6.0, shape=(), dtype=float32) Output = tf.Tensor(6.0, shape=(), dtype=float32) Output = tf.Tensor(6.0, shape=(), dtype=float32) Gradients and Jacobians: tf.GradientTape TensorFlow's GradientTape can be used for backpropagation through TensorFlow and TensorFlow NumPy code. Use the model created in Example Model section, and compute gradients and jacobians. Example Model def create_batch(batch_size=32): """Creates a batch of input and labels.""" return (tnp.random.randn(batch_size, 32).astype(tnp.float32), tnp.random.randn(batch_size, 2).astype(tnp.float32)) def compute_gradients(model, inputs, labels): """Computes gradients of squared loss between model prediction and labels.""" with tf.GradientTape() as tape: assert model.weights is not None # Note that `model.weights` need to be explicitly watched since they # are not tf.Variables. tape.watch(model.weights) # Compute prediction and loss prediction = model.predict(inputs) loss = tnp.sum(tnp.square(prediction - labels)) # This call computes the gradient through the computation above. return tape.gradient(loss, model.weights) inputs, labels = create_batch() gradients = compute_gradients(model, inputs, labels) # Inspect the shapes of returned gradients to verify they match the # parameter shapes. print("Parameter shapes:", [w.shape for w in model.weights]) print("Gradient shapes:", [g.shape for g in gradients]) # Verify that gradients are of type ND array. assert isinstance(gradients[0], tnp.ndarray) def create_batch(batch_size=32): """Creates a batch of input and labels.""" return (tnp.random.randn(batch_size, 32).astype(tnp.float32), tnp.random.randn(batch_size, 2).astype(tnp.float32)) def compute_gradients(model, inputs, labels): """Computes gradients of squared loss between model prediction and labels.""" with tf.GradientTape() as tape: assert model.weights is not None # Note that `model.weights` need to be explicitly watched since they # are not tf.Variables. tape.watch(model.weights) # Compute prediction and loss prediction = model.predict(inputs) loss = tnp.sum(tnp.square(prediction - labels)) # This call computes the gradient through the computation above. return tape.gradient(loss, model.weights) inputs, labels = create_batch() gradients = compute_gradients(model, inputs, labels) # Inspect the shapes of returned gradients to verify they match the # parameter shapes. print("Parameter shapes:", [w.shape for w in model.weights]) print("Gradient shapes:", [g.shape for g in gradients]) # Verify that gradients are of type ND array. assert isinstance(gradients[0], tnp.ndarray) Parameter shapes: [TensorShape([32, 64]), TensorShape([64]), TensorShape([64, 2])] Gradient shapes: [TensorShape([32, 64]), TensorShape([64]), TensorShape([64, 2])] Parameter shapes: [TensorShape([32, 64]), TensorShape([64]), TensorShape([64, 2])] Gradient shapes: [TensorShape([32, 64]), TensorShape([64]), TensorShape([64, 2])] # Computes a batch of jacobians. Each row is the jacobian of an element in the # batch of outputs w.r.t. the corresponding input batch element. def prediction_batch_jacobian(inputs): with tf.GradientTape() as tape: tape.watch(inputs) prediction = model.predict(inputs) return prediction, tape.batch_jacobian(prediction, inputs) inp_batch = tnp.ones([16, 32], tnp.float32) output, batch_jacobian = prediction_batch_jacobian(inp_batch) # Note how the batch jacobian shape relates to the input and output shapes. print("Output shape: %s, input shape: %s" % (output.shape, inp_batch.shape)) print("Batch jacobian shape:", batch_jacobian.shape) # Computes a batch of jacobians. Each row is the jacobian of an element in the # batch of outputs w.r.t. the corresponding input batch element. def prediction_batch_jacobian(inputs): with tf.GradientTape() as tape: tape.watch(inputs) prediction = model.predict(inputs) return prediction, tape.batch_jacobian(prediction, inputs) inp_batch = tnp.ones([16, 32], tnp.float32) output, batch_jacobian = prediction_batch_jacobian(inp_batch) # Note how the batch jacobian shape relates to the input and output shapes. print("Output shape: %s, input shape: %s" % (output.shape, inp_batch.shape)) print("Batch jacobian shape:", batch_jacobian.shape) Output shape: (16, 2), input shape: (16, 32) Batch jacobian shape: (16, 2, 32) Output shape: (16, 2), input shape: (16, 32) Batch jacobian shape: (16, 2, 32) Trace compilation: tf.function TensorFlow's tf.function works by "trace compiling" the code and then optimizing these traces for much faster performance. See the Introduction to Graphs and Functions. tf.function Introduction to Graphs and Functions tf.function can be used to optimize TensorFlow NumPy code as well. Here is a simple example to demonstrate the speedups. Note that the body of tf.function code includes calls to TensorFlow NumPy APIs. tf.function tf.function inputs, labels = create_batch(512) print("Eager performance") compute_gradients(model, inputs, labels) print(timeit.timeit(lambda: compute_gradients(model, inputs, labels), number=10) * 100, "ms") print("\ntf.function compiled performance") compiled_compute_gradients = tf.function(compute_gradients) compiled_compute_gradients(model, inputs, labels) # warmup print(timeit.timeit(lambda: compiled_compute_gradients(model, inputs, labels), number=10) * 100, "ms") inputs, labels = create_batch(512) print("Eager performance") compute_gradients(model, inputs, labels) print(timeit.timeit(lambda: compute_gradients(model, inputs, labels), number=10) * 100, "ms") print("\ntf.function compiled performance") compiled_compute_gradients = tf.function(compute_gradients) compiled_compute_gradients(model, inputs, labels) # warmup print(timeit.timeit(lambda: compiled_compute_gradients(model, inputs, labels), number=10) * 100, "ms") Eager performance 2.710705300000882 ms tf.function compiled performance 0.7041131000050882 ms Eager performance 2.710705300000882 ms tf.function compiled performance 0.7041131000050882 ms Vectorization: tf.vectorized_map TensorFlow has inbuilt support for vectorizing parallel loops, which allows speedups of one to two orders of magnitude. These speedups are accessible via the tf.vectorized_map API and apply to TensorFlow NumPy code as well. tf.vectorized_map It is sometimes useful to compute the gradient of each output in a batch w.r.t. the corresponding input batch element. Such computation can be done efficiently using tf.vectorized_map as shown below. tf.vectorized_map @tf.function def vectorized_per_example_gradients(inputs, labels): def single_example_gradient(arg): inp, label = arg return compute_gradients(model, tnp.expand_dims(inp, 0), tnp.expand_dims(label, 0)) # Note that a call to `tf.vectorized_map` semantically maps # `single_example_gradient` over each row of `inputs` and `labels`. # The interface is similar to `tf.map_fn`. # The underlying machinery vectorizes away this map loop which gives # nice speedups. return tf.vectorized_map(single_example_gradient, (inputs, labels)) batch_size = 128 inputs, labels = create_batch(batch_size) per_example_gradients = vectorized_per_example_gradients(inputs, labels) for w, p in zip(model.weights, per_example_gradients): print("Weight shape: %s, batch size: %s, per example gradient shape: %s " % ( w.shape, batch_size, p.shape)) @tf.function def vectorized_per_example_gradients(inputs, labels): def single_example_gradient(arg): inp, label = arg return compute_gradients(model, tnp.expand_dims(inp, 0), tnp.expand_dims(label, 0)) # Note that a call to `tf.vectorized_map` semantically maps # `single_example_gradient` over each row of `inputs` and `labels`. # The interface is similar to `tf.map_fn`. # The underlying machinery vectorizes away this map loop which gives # nice speedups. return tf.vectorized_map(single_example_gradient, (inputs, labels)) batch_size = 128 inputs, labels = create_batch(batch_size) per_example_gradients = vectorized_per_example_gradients(inputs, labels) for w, p in zip(model.weights, per_example_gradients): print("Weight shape: %s, batch size: %s, per example gradient shape: %s " % ( w.shape, batch_size, p.shape)) Weight shape: (32, 64), batch size: 128, per example gradient shape: (128, 32, 64) Weight shape: (64,), batch size: 128, per example gradient shape: (128, 64) Weight shape: (64, 2), batch size: 128, per example gradient shape: (128, 64, 2) Weight shape: (32, 64), batch size: 128, per example gradient shape: (128, 32, 64) Weight shape: (64,), batch size: 128, per example gradient shape: (128, 64) Weight shape: (64, 2), batch size: 128, per example gradient shape: (128, 64, 2) # Benchmark the vectorized computation above and compare with # unvectorized sequential computation using `tf.map_fn`. @tf.function def unvectorized_per_example_gradients(inputs, labels): def single_example_gradient(arg): inp, label = arg return compute_gradients(model, tnp.expand_dims(inp, 0), tnp.expand_dims(label, 0)) return tf.map_fn(single_example_gradient, (inputs, labels), fn_output_signature=(tf.float32, tf.float32, tf.float32)) print("Running vectorized computation") print(timeit.timeit(lambda: vectorized_per_example_gradients(inputs, labels), number=10) * 100, "ms") print("\nRunning unvectorized computation") per_example_gradients = unvectorized_per_example_gradients(inputs, labels) print(timeit.timeit(lambda: unvectorized_per_example_gradients(inputs, labels), number=10) * 100, "ms") # Benchmark the vectorized computation above and compare with # unvectorized sequential computation using `tf.map_fn`. @tf.function def unvectorized_per_example_gradients(inputs, labels): def single_example_gradient(arg): inp, label = arg return compute_gradients(model, tnp.expand_dims(inp, 0), tnp.expand_dims(label, 0)) return tf.map_fn(single_example_gradient, (inputs, labels), fn_output_signature=(tf.float32, tf.float32, tf.float32)) print("Running vectorized computation") print(timeit.timeit(lambda: vectorized_per_example_gradients(inputs, labels), number=10) * 100, "ms") print("\nRunning unvectorized computation") per_example_gradients = unvectorized_per_example_gradients(inputs, labels) print(timeit.timeit(lambda: unvectorized_per_example_gradients(inputs, labels), number=10) * 100, "ms") Running vectorized computation 0.659675699989748 ms Running unvectorized computation 29.259711299982882 ms Running vectorized computation 0.659675699989748 ms Running unvectorized computation 29.259711299982882 ms Device placement TensorFlow NumPy can place operations on CPUs, GPUs, TPUs and remote devices. It uses standard TensorFlow mechanisms for device placement. Below a simple example shows how to list all devices and then place some computation on a particular device. TensorFlow also has APIs for replicating computation across devices and performing collective reductions which will not be covered here. List devices tf.config.list_logical_devices and tf.config.list_physical_devices can be used to find what devices to use. tf.config.list_logical_devices tf.config.list_physical_devices print("All logical devices:", tf.config.list_logical_devices()) print("All physical devices:", tf.config.list_physical_devices()) # Try to get the GPU device. If unavailable, fallback to CPU. try: device = tf.config.list_logical_devices(device_type="GPU")[0] except IndexError: device = "/device:CPU:0" print("All logical devices:", tf.config.list_logical_devices()) print("All physical devices:", tf.config.list_physical_devices()) # Try to get the GPU device. If unavailable, fallback to CPU. try: device = tf.config.list_logical_devices(device_type="GPU")[0] except IndexError: device = "/device:CPU:0" All logical devices: [LogicalDevice(name='/device:CPU:0', device_type='CPU'), LogicalDevice(name='/device:GPU:0', device_type='GPU'), LogicalDevice(name='/device:GPU:1', device_type='GPU'), LogicalDevice(name='/device:GPU:2', device_type='GPU'), LogicalDevice(name='/device:GPU:3', device_type='GPU')] All physical devices: [PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'), PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU'), PhysicalDevice(name='/physical_device:GPU:1', device_type='GPU'), PhysicalDevice(name='/physical_device:GPU:2', device_type='GPU'), PhysicalDevice(name='/physical_device:GPU:3', device_type='GPU')] All logical devices: [LogicalDevice(name='/device:CPU:0', device_type='CPU'), LogicalDevice(name='/device:GPU:0', device_type='GPU'), LogicalDevice(name='/device:GPU:1', device_type='GPU'), LogicalDevice(name='/device:GPU:2', device_type='GPU'), LogicalDevice(name='/device:GPU:3', device_type='GPU')] All physical devices: [PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'), PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU'), PhysicalDevice(name='/physical_device:GPU:1', device_type='GPU'), PhysicalDevice(name='/physical_device:GPU:2', device_type='GPU'), PhysicalDevice(name='/physical_device:GPU:3', device_type='GPU')] Placing operations: tf.device tf.device Operations can be placed on a device by calling it in a tf.device scope. tf.device print("Using device: %s" % str(device)) # Run operations in the `tf.device` scope. # If a GPU is available, these operations execute on the GPU and outputs are # placed on the GPU memory. with tf.device(device): prediction = model.predict(create_batch(5)[0]) print("prediction is placed on %s" % prediction.device) print("Using device: %s" % str(device)) # Run operations in the `tf.device` scope. # If a GPU is available, these operations execute on the GPU and outputs are # placed on the GPU memory. with tf.device(device): prediction = model.predict(create_batch(5)[0]) print("prediction is placed on %s" % prediction.device) Using device: LogicalDevice(name='/device:GPU:0', device_type='GPU') prediction is placed on /job:localhost/replica:0/task:0/device:GPU:0 Using device: LogicalDevice(name='/device:GPU:0', device_type='GPU') prediction is placed on /job:localhost/replica:0/task:0/device:GPU:0 Copying ND arrays across devices: tnp.copy tnp.copy A call to tnp.copy, placed in a certain device scope, will copy the data to that device, unless the data is already on that device. tnp.copy with tf.device("/device:CPU:0"): prediction_cpu = tnp.copy(prediction) print(prediction.device) print(prediction_cpu.device) with tf.device("/device:CPU:0"): prediction_cpu = tnp.copy(prediction) print(prediction.device) print(prediction_cpu.device) /job:localhost/replica:0/task:0/device:GPU:0 /job:localhost/replica:0/task:0/device:CPU:0 /job:localhost/replica:0/task:0/device:GPU:0 /job:localhost/replica:0/task:0/device:CPU:0 Performance comparisons TensorFlow NumPy uses highly optimized TensorFlow kernels that can be dispatched on CPUs, GPUs and TPUs. TensorFlow also performs many compiler optimizations, like operation fusion, which translate to performance and memory improvements. See TensorFlow graph optimization with Grappler to learn more. TensorFlow graph optimization with Grappler However TensorFlow has higher overheads for dispatching operations compared to NumPy. For workloads composed of small operations (less than about 10 microseconds), these overheads can dominate the runtime and NumPy could provide better performance. For other cases, TensorFlow should generally provide better performance. Run the benchmark below to compare NumPy and TensorFlow NumPy performance for different input sizes. def benchmark(f, inputs, number=30, force_gpu_sync=False): """Utility to benchmark `f` on each value in `inputs`.""" times = [] for inp in inputs: def _g(): if force_gpu_sync: one = tnp.asarray(1) f(inp) if force_gpu_sync: with tf.device("CPU:0"): tnp.copy(one) # Force a sync for GPU case _g() # warmup t = timeit.timeit(_g, number=number) times.append(t * 1000. / number) return times def plot(np_times, tnp_times, compiled_tnp_times, has_gpu, tnp_times_gpu): """Plot the different runtimes.""" plt.xlabel("size") plt.ylabel("time (ms)") plt.title("Sigmoid benchmark: TF NumPy vs NumPy") plt.plot(sizes, np_times, label="NumPy") plt.plot(sizes, tnp_times, label="TF NumPy (CPU)") plt.plot(sizes, compiled_tnp_times, label="Compiled TF NumPy (CPU)") if has_gpu: plt.plot(sizes, tnp_times_gpu, label="TF NumPy (GPU)") plt.legend() def benchmark(f, inputs, number=30, force_gpu_sync=False): """Utility to benchmark `f` on each value in `inputs`.""" times = [] for inp in inputs: def _g(): if force_gpu_sync: one = tnp.asarray(1) f(inp) if force_gpu_sync: with tf.device("CPU:0"): tnp.copy(one) # Force a sync for GPU case _g() # warmup t = timeit.timeit(_g, number=number) times.append(t * 1000. / number) return times def plot(np_times, tnp_times, compiled_tnp_times, has_gpu, tnp_times_gpu): """Plot the different runtimes.""" plt.xlabel("size") plt.ylabel("time (ms)") plt.title("Sigmoid benchmark: TF NumPy vs NumPy") plt.plot(sizes, np_times, label="NumPy") plt.plot(sizes, tnp_times, label="TF NumPy (CPU)") plt.plot(sizes, compiled_tnp_times, label="Compiled TF NumPy (CPU)") if has_gpu: plt.plot(sizes, tnp_times_gpu, label="TF NumPy (GPU)") plt.legend() # Define a simple implementation of `sigmoid`, and benchmark it using # NumPy and TensorFlow NumPy for different input sizes. def np_sigmoid(y): return 1. / (1. + np.exp(-y)) def tnp_sigmoid(y): return 1. / (1. + tnp.exp(-y)) @tf.function def compiled_tnp_sigmoid(y): return tnp_sigmoid(y) sizes = (2 ** 0, 2 ** 5, 2 ** 10, 2 ** 15, 2 ** 20) np_inputs = [np.random.randn(size).astype(np.float32) for size in sizes] np_times = benchmark(np_sigmoid, np_inputs) with tf.device("/device:CPU:0"): tnp_inputs = [tnp.random.randn(size).astype(np.float32) for size in sizes] tnp_times = benchmark(tnp_sigmoid, tnp_inputs) compiled_tnp_times = benchmark(compiled_tnp_sigmoid, tnp_inputs) has_gpu = len(tf.config.list_logical_devices("GPU")) if has_gpu: with tf.device("/device:GPU:0"): tnp_inputs = [tnp.random.randn(size).astype(np.float32) for size in sizes] tnp_times_gpu = benchmark(compiled_tnp_sigmoid, tnp_inputs, 100, True) else: tnp_times_gpu = None plot(np_times, tnp_times, compiled_tnp_times, has_gpu, tnp_times_gpu) # Define a simple implementation of `sigmoid`, and benchmark it using # NumPy and TensorFlow NumPy for different input sizes. def np_sigmoid(y): return 1. / (1. + np.exp(-y)) def tnp_sigmoid(y): return 1. / (1. + tnp.exp(-y)) @tf.function def compiled_tnp_sigmoid(y): return tnp_sigmoid(y) sizes = (2 ** 0, 2 ** 5, 2 ** 10, 2 ** 15, 2 ** 20) np_inputs = [np.random.randn(size).astype(np.float32) for size in sizes] np_times = benchmark(np_sigmoid, np_inputs) with tf.device("/device:CPU:0"): tnp_inputs = [tnp.random.randn(size).astype(np.float32) for size in sizes] tnp_times = benchmark(tnp_sigmoid, tnp_inputs) compiled_tnp_times = benchmark(compiled_tnp_sigmoid, tnp_inputs) has_gpu = len(tf.config.list_logical_devices("GPU")) if has_gpu: with tf.device("/device:GPU:0"): tnp_inputs = [tnp.random.randn(size).astype(np.float32) for size in sizes] tnp_times_gpu = benchmark(compiled_tnp_sigmoid, tnp_inputs, 100, True) else: tnp_times_gpu = None plot(np_times, tnp_times, compiled_tnp_times, has_gpu, tnp_times_gpu) Further reading TensorFlow NumPy: Distributed Image Classification Tutorial TensorFlow NumPy: Keras and Distribution Strategy Sentiment Analysis with Trax and TensorFlow NumPy TensorFlow NumPy: Distributed Image Classification Tutorial TensorFlow NumPy: Distributed Image Classification Tutorial TensorFlow NumPy: Keras and Distribution Strategy TensorFlow NumPy: Keras and Distribution Strategy Sentiment Analysis with Trax and TensorFlow NumPy Sentiment Analysis with Trax and TensorFlow NumPy Originally published on the TensorFlow website, this article appears here under a new headline and is licensed under CC BY 4.0. Code samples shared under the Apache 2.0 License. Originally published on the TensorFlow website, this article appears here under a new headline and is licensed under CC BY 4.0. Code samples shared under the Apache 2.0 License. Originally published on the TensorFlow website, this article appears here under a new headline and is licensed under CC BY 4.0. Code samples shared under the Apache 2.0 License. TensorFlow TensorFlow