Spaces:
Running
Running
| # Lint as: python3 | |
| # Copyright 2020 The TensorFlow Authors. All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # ============================================================================== | |
| """Functions and classes related to training performance.""" | |
| import tensorflow as tf | |
| def configure_optimizer(optimizer, | |
| use_float16=False, | |
| use_graph_rewrite=False, | |
| loss_scale="dynamic"): | |
| """Configures optimizer object with performance options.""" | |
| if use_float16: | |
| # Wraps optimizer with a LossScaleOptimizer. This is done automatically | |
| # in compile() with the "mixed_float16" policy, but since we do not call | |
| # compile(), we must wrap the optimizer manually. | |
| optimizer = ( | |
| tf.keras.mixed_precision.experimental.LossScaleOptimizer( | |
| optimizer, loss_scale=loss_scale)) | |
| if use_graph_rewrite: | |
| # Note: the model dtype must be 'float32', which will ensure | |
| # tf.ckeras.mixed_precision and | |
| # tf.train.experimental.enable_mixed_precision_graph_rewrite do not double | |
| # up. | |
| optimizer = tf.train.experimental.enable_mixed_precision_graph_rewrite( | |
| optimizer) | |
| return optimizer | |
| def set_mixed_precision_policy(dtype, loss_scale=None): | |
| """Sets mix precision policy.""" | |
| if dtype == tf.float16: | |
| policy = tf.keras.mixed_precision.experimental.Policy( | |
| 'mixed_float16', loss_scale=loss_scale) | |
| tf.keras.mixed_precision.experimental.set_policy(policy) | |
| elif dtype == tf.bfloat16: | |
| policy = tf.keras.mixed_precision.experimental.Policy( | |
| 'mixed_bfloat16') | |
| tf.keras.mixed_precision.experimental.set_policy(policy) | |
| elif dtype == tf.float32: | |
| tf.keras.mixed_precision.experimental.set_policy('float32') | |
| else: | |
| raise ValueError("Unexpected dtype: %s" % dtype) | |