Is there an easy way to get something like Keras model.summary in Tensorflow?

Viewed 37476

I have been working with Keras and really liked the model.summary() It gives a good overview of the size of the different layers and especially an overview of the number of parameters the model has.

Is there a similar function in Tensorflow? I could find nothing on Stackoverflow or the Tensorflow API documentation.

3 Answers

Looks like you can use Slim

Example:

import numpy as np

from tensorflow.python.layers import base
import tensorflow as tf
import tensorflow.contrib.slim as slim

x = np.zeros((1,4,4,3))
x_tf = tf.convert_to_tensor(x, np.float32)
z_tf = tf.layers.conv2d(x_tf, filters=32, kernel_size=(3,3))

def model_summary():
    model_vars = tf.trainable_variables()
    slim.model_analyzer.analyze_vars(model_vars, print_info=True)

model_summary()

Output:

---------
Variables: name (type shape) [size]
---------
conv2d/kernel:0 (float32_ref 3x3x3x32) [864, bytes: 3456]
conv2d/bias:0 (float32_ref 32) [32, bytes: 128]
Total size of variables: 896
Total bytes of variables: 3584

Also here is an example of custom function to print model summary: https://github.com/NVlabs/stylegan/blob/f3a044621e2ab802d40940c16cc86042ae87e100/dnnlib/tflib/network.py#L507

If you already have .pb tensorflow model you can use: inspect_pb.py to print model info or use tensorflow summarize_graph tool with --print_structure flag, also it's nice that it can detect input and output names.

You can use keras with the tensorflow backend to get the best features of either keras or tensorflow.

Related