search

galeone / dynamic-training-bench

Simplify the training and tuning of Tensorflow models
Mozilla Public License 2.0
213 stars 31 forks source link

how to compute the GPU memory that the graph need be using ? #12

Closed guiyang882 closed 7 years ago

guiyang882 commented 7 years ago

hi, I found a question, when i want to dispatch my job to different machines. I can know the machines ability, but how to estimator the graph's GPU memory ? Please give me some idea, thanks a lot .

thx !

galeone commented 7 years ago

In the latest version of DyTB when you start a training you'll see the model architecture and a short summary of the model size (number of parameters and size in KB):

Model : trainable parameters: . Size: KB

If you just want to know the size of the model you have to define the model first and the use count_trainable_parameters, something like this (not tested):

from dytb.models.utils import count_trainable_parameters
# import your model here, lets call it yourmodel
# define a placeholder called inputs with the expected size of the network input
model = yourmodel.get(inputs, train_phase=True)
num_of_parameters = count_trainable_parameters(print_model=True)
print("Num of parameters {}, size KB: {}".format(num_of_parameters, num_of_parameters * 4/1000))
guiyang882 commented 7 years ago

oh, i found it. that code can help me to determine the model size, and I can estimator the GPU memory size. I will try it, if I have some problem, I will border you for helping !☕️☕️☕️

galeone commented 7 years ago

No problem!