How to count total number of trainable parameters in a tensorflow model?

Question

Is there a function call or another way to count the total number of parameters in a tensorflow model?

By parameters I mean: an N dim vector of trainable variables has N parameters, a NxM matrix has N*M parameters, etc. So essentially I'd like to sum the product of the shape dimensions of all the trainable variables in a tensorflow session.

your question description and title do not match (unless I'm confusing the terminology of graph and model). In the question you ask about a graph and the title you ask about a model. What if you had two different models? I'd suggest to clarify that on the question. — Charlie Parker, Nov 21 '16 at 23:39
Related if you are using Keras: https://stackoverflow.com/questions/45046525/how-can-i-get-the-number-of-trainable-parameters-of-a-model-in-keras — bers, Jun 08 '20 at 20:10

nessuno · Accepted Answer · 2018-12-06T07:54:39.353

93

Loop over the shape of every variable in tf.trainable_variables().

total_parameters = 0
for variable in tf.trainable_variables():
    # shape is an array of tf.Dimension
    shape = variable.get_shape()
    print(shape)
    print(len(shape))
    variable_parameters = 1
    for dim in shape:
        print(dim)
        variable_parameters *= dim.value
    print(variable_parameters)
    total_parameters += variable_parameters
print(total_parameters)

Update: I wrote an article to clarify the dynamic/static shapes in Tensorflow because of this answer: https://pgaleone.eu/tensorflow/2018/07/28/understanding-tensorflow-tensors-shape-static-dynamic/

edited Dec 06 '18 at 07:54

answered Jul 02 '16 at 14:42

nessuno

26,493
5
83
74

5

if you have more than one model, how does `tf.trainable_variables()` know which one to use? – Charlie Parker Nov 21 '16 at 23:33
2

tf.trainable_variables() returns all the variables marked as trainable that are present in the current graph. If in the current graph you have more than one model, you have to manually filter the variables using theyr names. Somethink like if variable.name.strartswith("model2"): ... – nessuno Nov 22 '16 at 06:54
this solution gives me the error "Exception occurred: Can't convert 'int' object to str implicitly". You need to cast 'dim' explicitly to 'int' as the suggested in the answer below (which I suggest to be the correct answer) – whiletrue Oct 07 '17 at 14:51
really helpful, – Sudip Das Apr 02 '18 at 10:17
It seems in TF2, this has changed to tf.compat.v1.trainable_variables()! But then this return 0 parameters! – chikitin Nov 11 '19 at 14:16
3

In TensorFlow 2 this answer is deprecated. You have to use the `.trainable_variables` of your Keras model - there is no more a global graph! – nessuno Nov 11 '19 at 19:22

Michael Gygli · Answer 2 · 2017-05-06T11:34:16.283

49

I have an even shorter version, one line solution using using numpy:

np.sum([np.prod(v.get_shape().as_list()) for v in tf.trainable_variables()])

edited May 06 '17 at 11:34

answered May 05 '17 at 09:15

Michael Gygli

900
6
11

in my version, v doesn't have a shape_as_list() function but only get_shape() function – mustafa May 06 '17 at 00:25
I think earlier versions don't have .shape but get_shape(). Updated my answer. Anyway, I wrote v.shape.as_list() and not v.shape_as_list(). – Michael Gygli May 06 '17 at 11:33
13

`np.sum([np.prod(v.shape) for v in tf.trainable_variables()])` works as well in TensorFlow 1.2 – Julius Kunze Jul 02 '17 at 15:51
1

`np.sum([np.prod(v.shape) for v in model.trainable_variables])` works for me --> without function call "()" at the end – Manute Sep 14 '22 at 08:52

score 9 · Answer 3 · edited Nov 21 '16 at 23:34

9

Not sure if the answer given actually runs (I found you need to convert the dim object to an int for it to work). Here is is one that works and you can just copy paste the functions and call them (added a few comments too):

def count_number_trainable_params():
    '''
    Counts the number of trainable variables.
    '''
    tot_nb_params = 0
    for trainable_variable in tf.trainable_variables():
        shape = trainable_variable.get_shape() # e.g [D,F] or [W,H,C]
        current_nb_params = get_nb_params_shape(shape)
        tot_nb_params = tot_nb_params + current_nb_params
    return tot_nb_params

def get_nb_params_shape(shape):
    '''
    Computes the total number of params for a given shap.
    Works for any number of shapes etc [D,F] or [W,H,C] computes D*F and W*H*C.
    '''
    nb_params = 1
    for dim in shape:
        nb_params = nb_params*int(dim)
    return nb_params

edited Nov 21 '16 at 23:34

f4.

3,814
1
23
30

answered Oct 08 '16 at 20:16

Charlie Parker

5,884
57
198
323

the answer does work (r0.11.0). yours is more plug n play :) – f4. Nov 21 '16 at 23:34
@f4. there seems to be a bug with this because `y` doesn't seem to be used. – Charlie Parker Nov 21 '16 at 23:34
1

@CharlieParker I fixed it a few seconds ago ;) – f4. Nov 21 '16 at 23:35
@f4. it still doesn't truly solve the issue I was trying to do (or the original author intended since he gave y as an input) because I was looking for a function that depended on the model one gave as input (i.e. `y`). Right now as given, I have no idea what on earth it counts. My suspicion is that it counts just all models (I have two separate models). – Charlie Parker Nov 21 '16 at 23:36
@CharlieParker it counts all trainable variables, which by default is all variables I believe. You can work something out using the variables attributes like graph or name. – f4. Nov 21 '16 at 23:43

score 7 · Answer 4 · 2020-04-28T16:56:42.413

7

Update April 2020: tfprof and the Profiler UI have been deprecated in favor of profiler support in TensorBoard.

The two existing answers are good if you're looking into computing the number of parameters yourself. If your question was more along the lines of "is there an easy way to profile my TensorFlow models?", I would highly recommend looking into tfprof. It profiles your model, including calculating the number of parameters.

edited Apr 28 '20 at 16:56

answered Dec 21 '16 at 06:45

The `tfprof` link is broken. As the edit queue is full here is the [working link](https://github.com/tensorflow/profiler-ui). Also for the `tfprof` has been deprecated. – dsalaj Apr 25 '20 at 09:28

score 3 · Answer 5 · answered Dec 29 '16 at 15:00

I'll throw in my equivalent but shorter implementation:

def count_params():
    "print number of trainable variables"
    size = lambda v: reduce(lambda x, y: x*y, v.get_shape().as_list())
    n = sum(size(v) for v in tf.trainable_variables())
    print "Model size: %dK" % (n/1000,)

score 3 · Answer 6 · answered Aug 07 '18 at 02:44

If one prefers to avoid numpy (it can be left out for many projects), then:

all_trainable_vars = tf.reduce_sum([tf.reduce_prod(v.shape) for v in tf.trainable_variables()])

This is a TF translation of the previous answer by Julius Kunze.

As any TF operation, it requires a session run to evaluate:

print(sess.run(all_trainable_vars))

score 3 · Answer 7 · answered Dec 14 '21 at 07:11

3

Now, you can use this :

from keras.utils.layer_utils import count_params  

count_params(model.trainable_weights)

answered Dec 14 '21 at 07:11

jeugregg

116
6

score 2 · Answer 8 · answered Aug 19 '22 at 09:49

Works for me on TF v2.9. Credit to this answer

import numpy as np

trainable_params = np.sum([np.prod(v.get_shape()) for v in model.trainable_weights])
non_trainable_params = np.sum([np.prod(v.get_shape()) for v in model.non_trainable_weights])
total_params = trainable_params + non_trainable_params
    
print(trainable_params)
print(non_trainable_params)
print(total_params)

score -1 · Answer 9 · edited Jul 30 '23 at 06:36

model.summary()

Model: "sequential_32"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
conv2d_88 (Conv2D)           (None, 240, 240, 16)      448       
_________________________________________________________________
max_pooling2d_87 (MaxPooling (None, 120, 120, 16)      0         
_________________________________________________________________
conv2d_89 (Conv2D)           (None, 120, 120, 32)      4640      
_________________________________________________________________
max_pooling2d_88 (MaxPooling (None, 60, 60, 32)        0         
_________________________________________________________________
conv2d_90 (Conv2D)           (None, 60, 60, 64)        18496     
_________________________________________________________________
max_pooling2d_89 (MaxPooling (None, 30, 30, 64)        0         
_________________________________________________________________
flatten_29 (Flatten)         (None, 57600)             0         
_________________________________________________________________
dropout_48 (Dropout)         (None, 57600)             0         
_________________________________________________________________
dense_150 (Dense)            (None, 24)                1382424   
_________________________________________________________________
dense_151 (Dense)            (None, 9)                 225       
_________________________________________________________________
dense_152 (Dense)            (None, 3)                 30        
_________________________________________________________________
dense_153 (Dense)            (None, 1)                 4         
=================================================================
Total params: 1,406,267
Trainable params: 1,406,267
Non-trainable params: 0
_________________________________________________________________

This is Keras, not Tensorflow; question is clearly about Tensorflow models, not Keras ones. — desertnaut, Jun 19 '20 at 01:32

How to count total number of trainable parameters in a tensorflow model?

9 Answers9

Linked