def try_count_params(model: tf.keras.Model):
"""Count the number of parameters if model is possible.
Args:
model: Try to count the number of params in this model.
Returns:
The number of parameters or None.
"""
if hasattr(model, 'count_params'):
try:
return model.count_params()
except ValueError:
logging.info(
'Number of trainable params unknown, because the build() '
'methods in keras layers were not called. This is probably '
'because the model was not feed any input, e.g., the max '
'train step already reached before this run.')
return None
return None
def calc_max_memory_usage(model: tf.keras.Model):
"""Calculate memory requirement of a model per sample in bits."""
layers = extract_layers(model)
n_shapes = int(numpy.sum(
[numpy.prod(numpy.array([s if isinstance(s, int) else 1 for s in l.output_shape])) for l in layers]))
n_parameters = model.count_params()
# memory needed for saving activations during gradient calculation
n_activations = 0
for l in layers:
if len(l.trainable_variables) == 0 or l.output_shape is None:
continue
activation_shapes = l.output_shape
if not isinstance(activation_shapes[0], tuple):
activation_shapes = [tuple(activation_shapes)]
print(activation_shapes)
print(n_activations)
total_memory = (n_shapes + n_parameters + n_activations) * 32
return total_memory * 1.1641532182693481 * 10 ** -10
