How to check that Tensorflow graph rewrites that use MKL occur?

Question

From looking at Tensorflow code, some MKL optimizations are done by a graph rewrite replacing sets of nodes by fused functions that use MKL. I tried to look for the rewrites with tf.logging.set_verbosity(1) but never see of the log messages I expect.

I have built Tensorflow from sources on CPU with MKL and XLA enabled. I think the build is using MKL because I can use 'NCHW' data format for tf.nn.conv2d and tf.nn.bias_add in the forward pass if they occur together. It also runs faster and fully utilises the CPU. The backward pass though errors saying that "CPU BiasGradOp only supports NHWC" though it looks like MKL functions exist to fuse Conv2D and BiasAdd both forwards and backwards with 'NCHW'. So I want to look directly for the rewrites.

How can I see if the graph rewrites are happening?

Answer 1

One way is to use the timeline/trace feature. You can followed this StackOverflow answer. If it uses MKL you would see nodes with names like _MklReshape or _MklConv2D

Answer 2

This is not specifically testing for graph rewrites, but you can check, if mkl is enabled in tensorflow by using:

tf.python.pywrap_tensorflow.IsMklEnabled()

From: https://github.com/tensorflow/tensorflow/issues/17176#issuecomment-371364155

Answer 3

Tensorflow has a debugger (tfdbg) with a tutorial here. The debugger prints a list of all graph nodes that will be visited by a session.run() before running it.

You can also explore the input tensors, output tensors, and the attributes of each node.

Ariel's answer also works to see the op types if you don't want to take the time to compile with tfdbg.

Answer 4

For v2.0.0+ the command is:

python -c "from tensorflow.python import pywrap_tensorflow; print(pywrap_tensorflow.IsMklEnabled())"

source: https://software.intel.com/en-us/forums/intel-optimized-ai-frameworks/topic/837000