tf.function compiles a Python function into an executable TensorFlow program, usually in graph mode. The compiled version can run faster, but the executed code differs from the Python you wrote, making debugging harder. Debug the function first, then enable compilation. section 3.3.1
Compile a dense layer
@tf.function def dense(inputs, W, b): return tf.nn.relu(tf.matmul(inputs, W) + b) The first call takes longer because TensorFlow compiles the function. Later calls reuse the compiled program and are faster. section 3.3.1
Use these practical rules when applying tf.function:
- Debug in ordinary Python execution before decorating the function. section 3.3.1
- Expect compilation overhead on the first call and faster subsequent calls. section 3.3.1
- Use the default graph mode for standard compilation, or request XLA with jit_compile=True when additional optimization is useful. section 3.3.1
@tf.function def dense(inputs, W, b): return tf.nn.relu(tf.matmul(inputs, W) + b)
Inspect the compact decorator example for a compiled TensorFlow dense-layer function. section 3.3.1
@tf.function(jit_compile=True) def dense(inputs, W, b): return tf.nn.relu(tf.matmul(inputs, W) + b)
Compare the default decorator with XLA compilation enabled through jit_compile=True. section 3.3.1
