experiments
/
STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Allow tf.distribute.TPUStrategy to be used with TPUEmbedding API and ensure that LossScaleOptimizer properly rejects it. 

PiperOrigin-RevId: 318186211
Change-Id: Id3b9cb8288e5d28ddbaec97d5b35627ab35bc08d
This commit is contained in:
Bruce Fontaine 2020-06-24 19:00:55 -07:00 committed by Geeta Chavan
parent 890eae3e88
commit b878ae3918
2 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/python/keras/mixed_precision/experimental/loss_scale_optimizer.py
View File

@ -440,7 +440,8 @@ class LossScaleOptimizer(_DelegatingTrackableMixin, optimizer_v2.OptimizerV2):
if not strategy_supports_loss_scaling():
strategy = distribution_strategy_context.get_strategy()
if isinstance(strategy,
(tpu_strategy.TPUStrategy, tpu_strategy.TPUStrategyV1)):
(tpu_strategy.TPUStrategy, tpu_strategy.TPUStrategyV1,
tpu_strategy.TPUStrategyV2)):
raise ValueError(
'Loss scaling is not supported with TPUStrategy. Loss scaling is '
'unnecessary with TPUs, since they support bfloat16 instead of '

3
tensorflow/python/tpu/tpu_embedding_v2.py
View File

@ -265,7 +265,8 @@ class TPUEmbedding(tracking.AutoTrackable):
Adam or Adagrad).
"""
self._strategy = distribution_strategy_context.get_strategy()
self._using_tpu = isinstance(self._strategy, tpu_strategy.TPUStrategy)
self._using_tpu = isinstance(self._strategy, (tpu_strategy.TPUStrategy,
tpu_strategy.TPUStrategyV2))
self._pipeline_execution_with_tensor_core = (
pipeline_execution_with_tensor_core)