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