-
Notifications
You must be signed in to change notification settings - Fork 3.6k
Refine automatic mixed precision support via hyper param #1681
Conversation
tensor2tensor/utils/optimize.py
Outdated
@@ -71,8 +69,7 @@ def optimize(loss, | |||
opt = ConditionalOptimizer(hparams.optimizer, learning_rate, hparams, use_tpu) | |||
if use_tpu: | |||
opt = tf.contrib.tpu.CrossShardOptimizer(opt) | |||
if gpu_auto_mixed_precision or os.environ.get( | |||
"TF_ENABLE_AUTO_MIXED_PRECISION", "0") == "1": | |||
if hparams.gpu_automatic_mixed_precision: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
get(hparams, "gpu_automatic_mixed_precision", False) is preferable -- since people may pass an hparam that doesn't have this param -- for example in tests etc.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
good one. I fixed this
memory_height=1 | ||
memory_height=1, | ||
# Whether to use GPU automatic mixed precision (via graph rewrite) | ||
gpu_automatic_mixed_precision=False |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is good, but as in your earlier PR, based on a flag can you set this to true?
i.e. after we make the hparams in t2t_trainer, based on your flag, flip this on
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Just add the flag again to trainer and turn hparams on accordingly.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
done
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Left a few comments -- thanks for the changes
Thanks for the feedbacks @afrozenator . Let me know if the latest revision works. |
Thanks a lot @vinhngx for contributing this in the first place and now making it better! Will merge it in shortly. |
great thanks. I'm closing #1680 then. |
PiperOrigin-RevId: 266390503
In continuation of #1637 and in response to @afrozenator 's comments in #1680
In this PR, we re-organize automatic mixed precision training support to provide a cleaner implementation and an easier interface via using hyper parameters.
In particular, GPU automatic mixed precision training can now be enabled via setting a flag (and correspondingly a so-named hyper-parameter)
gpu_automatic_mixed_precision
for all tensor2tensor models, for example:Transformer
Resnet:
This is opposed to the previous approaches of setting the OS flag
TF_ENABLE_AUTO_MIXED_PRECISION
which is a non-programatic approach, or passing the flaggpu_auto_mixed_precision
directly to the optimizer (which will require modification of individual models to make call to optimizer with mixed precision training option).