Skip to content

GitLab

Commit 82f46dc7 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower
Browse files

Internal change 

PiperOrigin-RevId: 342098346
parent 865631be
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
official/core/train_lib.py
View file @ 82f46dc7
...@@ -228,6 +228,9 @@ def run_experiment(distribution_strategy: tf.distribute.Strategy, ...@@ -228,6 +228,9 @@ def run_experiment(distribution_strategy: tf.distribute.Strategy,
else: else:
raise NotImplementedError('The mode is not implemented: %s' % mode) raise NotImplementedError('The mode is not implemented: %s' % mode)
if hasattr(trainer.model, 'count_params'):
logging.info('Number of trainable params in model: %f Millions.',
trainer.model.count_params() / 10.**6)
if run_post_eval: if run_post_eval:
with distribution_strategy.scope(): with distribution_strategy.scope():
return trainer.model, trainer.evaluate( return trainer.model, trainer.evaluate(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment