def _enable_mixed_precision_graph_rewrite_base(opt, loss_scale, use_v1_behavior): """Enables mixed precision. See `enable_mixed_precision_graph_rewrite`.""" if not mixed_precision_global_state.using_default_mixed_precision_policy: raise ValueError( 'The mixed precision graph rewrite cannot be enabled, because a keras ' 'mixed precision Policy has been set. At most, one of the following ' 'functions can be called:\n\n' ' 1. tf.keras.mixed_precision.experimental.set_policy() (You called ' 'this first)\n' ' 2. tf.train.experimental.enable_mixed_precision_graph_rewrite() ' '(You called this second)\n\n' 'You called both functions, which is an error, because both functions ' 'enable you to use mixed precision. If in doubt which function to use, ' 'use the second, as it is currently more complete and easy to use. The ' 'second function enables mixed precision in the graph with a graph ' 'rewrite. However it is currently not very customizable, and does not ' 'support eager.') if mixed_precision_global_state.non_mixed_precision_session_created: # TODO(reedwm): Give the stacktrace of the existing Sessions. And if the # Sessions have already been closed, do not raise this error message. tf_logging.warn( 'You already have existing Sessions that do not use mixed ' 'precision. enable_mixed_precision_graph_rewrite() will ' 'not affect these Sessions.') opt = _wrap_optimizer(opt, loss_scale, use_v1_behavior=use_v1_behavior) config.set_optimizer_experimental_options({'auto_mixed_precision': True}) mixed_precision_global_state.mixed_precision_graph_rewrite_is_enabled = True return opt
def setUp(self): super(PlacementTest, self).setUp() # Grappler optimizations can affect whether the placement issues occur, # since they may inadvertently rewrite nodes and edges in a way that removes # cross-device copies. config.set_optimizer_experimental_options( {"disable_meta_optimizer": True})
def _enable_mixed_precision_graph_rewrite_base(opt, loss_scale, use_v1_behavior): """Enables mixed precision. See `enable_mixed_precision_graph_rewrite`.""" if mixed_precision_global_state.using_mixed_precision_policy: raise ValueError( 'The mixed precision graph rewrite cannot be enabled, because the ' 'global Keras dtype Policy has been set to a mixed precision policy. ' 'At most, one of the following can be called:\n\n' ' 1. tf.keras.mixed_precision.experimental.set_policy() with a mixed ' 'precision policy (You called this first)\n\n' ' 2. tf.train.experimental.enable_mixed_precision_graph_rewrite() ' '(You called this second)\n' 'You called both functions, which is an error, because both functions ' 'enable you to use mixed precision. If in doubt which function to use, ' 'use the first, as it supports Eager execution and is more ' 'customizable.') if mixed_precision_global_state.non_mixed_precision_session_created: # TODO(reedwm): Give the stacktrace of the existing Sessions. And if the # Sessions have already been closed, do not raise this error message. tf_logging.warn('You already have existing Sessions that do not use mixed ' 'precision. enable_mixed_precision_graph_rewrite() will ' 'not affect these Sessions.') opt = _wrap_optimizer(opt, loss_scale, use_v1_behavior=use_v1_behavior) config.set_optimizer_experimental_options({'auto_mixed_precision': True}) mixed_precision_global_state.mixed_precision_graph_rewrite_is_enabled = True return opt
def disable_mixed_precision_graph_rewrite_v1(): """Disables the mixed precision graph rewrite. After this is called, the mixed precision graph rewrite will no longer run for new Sessions, and so float32 operations will no longer be converted to float16 in such Sessions. However, any existing Sessions will continue to have the graph rewrite enabled if they were created after `enable_mixed_precision_graph_rewrite` was called but before `disable_mixed_precision_graph_rewrite` was called. This does not undo the effects of loss scaling. Any optimizers wrapped with a LossScaleOptimizer will continue to do loss scaling, although this loss scaling will no longer be useful if the optimizer is used in new Sessions, as the graph rewrite no longer converts the graph to use float16. This function is useful for unit testing. A unit tests can test using the mixed precision graph rewrite, then disable it so future unit tests continue using float32. If this is done, unit tests should not share a single session, as `enable_mixed_precision_graph_rewrite` and `disable_mixed_precision_graph_rewrite` have no effect on existing sessions. """ # We only have a separate V1 version of this function, because the V1 # docstring mentions sessions. if (not mixed_precision_global_state.is_mixed_precision_graph_rewrite_enabled()): tf_logging.warn('disable_mixed_precision_graph_rewrite() called when mixed ' 'precision is already disabled.') config.set_optimizer_experimental_options({'auto_mixed_precision': False}) mixed_precision_global_state.set_mixed_precision_graph_rewrite_enabled(False)
def tearDown(self): # Set auto_mixed_precision back to it's default value. config.set_optimizer_experimental_options( {'auto_mixed_precision': False}) # Set the IGNORE_PERF_VAR variable back to it's original value. if self._original_ignore_perf_value is not None: os.environ[self.IGNORE_PERF_VAR] = self._original_ignore_perf_value else: del os.environ[self.IGNORE_PERF_VAR] super(MixedPrecisionTest, self).tearDown()
def _enable_mixed_precision_graph_rewrite_base(opt, loss_scale, use_v1_behavior): """Enables mixed precision. See `enable_mixed_precision_graph_rewrite`.""" if mixed_precision_global_state.non_mixed_precision_session_created: # TODO(reedwm): Give the stacktrace of the existing Sessions. And if the # Sessions have already been closed, do not raise this error message. tf_logging.warn('You already have existing Sessions that do not use mixed ' 'precision. enable_mixed_precision_graph_rewrite() will ' 'not affect these Sessions.') opt = _wrap_optimizer(opt, loss_scale, use_v1_behavior=use_v1_behavior) config.set_optimizer_experimental_options({'auto_mixed_precision': True}) mixed_precision_global_state.mixed_precision_is_enabled = True return opt
def enable_mixed_precision_graph_rewrite(opt, loss_scale='dynamic'): """Enable mixed precision via a graph rewrite. Mixed precision is the use of both float16 and float32 when training a model, and is used to make the model run faster. This function will use mixed precision to speed up the execution time of your model when run on a GPU. It does this by changing the dtype of certain operations in the graph from float32 to float16. This function additionally wraps an Optimizer with a LossScaleOptimizer, which is required to prevent underflow in the float16 tensors during the backwards pass. An optimizer must be passed to this function, which will then be wrapped to use loss scaling. When this function is used, gradients should only be computed and applied with the returned optimizer, either by calling `opt.minimize()` or `opt.compute_gradients()` followed by `opt.apply_gradients()`. Gradients should not be computed with `tf.gradients` or `tf.GradientTape`. This is because the returned optimizer will apply loss scaling, and `tf.gradients`/`tf.GradientTape` will not. If you do directly use `tf.gradients` or `tf.GradientTape`, your model may train to a worse quality. When eager execution is enabled, the mixed precision graph rewrite is only enabled within `tf.function`s, as outside `tf.function`s, there is no graph. When enabled, mixed precision is only used on Volta GPUs and above. The parts of the graph on CPUs and TPUs are untouched by the graph rewrite. Args: opt: An instance of a `tf.keras.optimizers.Optimizer` or a `tf.train.Optimizer`. loss_scale: Either an int/float, the string "dynamic", or an instance of a `tf.train.experimental.LossScale`. The loss scale to use. It is recommended to keep this as its default value of "dynamic". Returns: A version of `opt` that will use loss scaling to prevent underflow. """ # TODO(reedwm): If a ConfigProto is passed to Session, either assert that # auto_mixed_precision is on or turn it on for the user. if mixed_precision_global_state.non_mixed_precision_session_created: # TODO(reedwm): Give the stacktrace of the existing Sessions. And if the # Sessions have already been closed, do not raise this error message. tf_logging.warn( 'You already have existing Sessions that do not use mixed ' 'precision. enable_mixed_precision_graph_rewrite() will ' 'not affect these Sessions.') opt = _wrap_optimizer(opt, loss_scale) config.set_optimizer_experimental_options({'auto_mixed_precision': True}) mixed_precision_global_state.mixed_precision_is_enabled = True return opt
def _enable_mixed_precision_graph_rewrite_base(opt, loss_scale, use_v1_behavior): """Enables mixed precision. See `enable_mixed_precision_graph_rewrite`.""" if mixed_precision_global_state.non_mixed_precision_session_created: # TODO(reedwm): Give the stacktrace of the existing Sessions. And if the # Sessions have already been closed, do not raise this error message. tf_logging.warn( 'You already have existing Sessions that do not use mixed ' 'precision. enable_mixed_precision_graph_rewrite() will ' 'not affect these Sessions.') opt = _wrap_optimizer(opt, loss_scale, use_v1_behavior=use_v1_behavior) config.set_optimizer_experimental_options({'auto_mixed_precision': True}) mixed_precision_global_state.mixed_precision_is_enabled = True return opt
def testOptimizerBoolOption(self, field): # TODO(b/128531235): Improve testing of option options = config.get_optimizer_experimental_options() self.assertFalse(options.get(field)) config.set_optimizer_experimental_options({field: True}) options[field] = True self.assertDictEqual(config.get_optimizer_experimental_options(), options) self.assertDictEqual(context.context().get_optimizer_experimental_options(), options) config.set_optimizer_experimental_options({field: False}) options[field] = False self.assertDictEqual(config.get_optimizer_experimental_options(), options) self.assertDictEqual(context.context().get_optimizer_experimental_options(), options)
def testOptimizerBoolOption(self, field): # TODO(b/128531235): Improve testing of option options = config.get_optimizer_experimental_options() self.assertFalse(options.get(field)) config.set_optimizer_experimental_options({field: True}) options[field] = True self.assertDictEqual(config.get_optimizer_experimental_options(), options) self.assertDictEqual( context.context().get_optimizer_experimental_options(), options) config.set_optimizer_experimental_options({field: False}) options[field] = False self.assertDictEqual(config.get_optimizer_experimental_options(), options) self.assertDictEqual( context.context().get_optimizer_experimental_options(), options)
def disable_mixed_precision_graph_rewrite(): """Disables the mixed precision graph rewrite. After this is called, the mixed precision graph rewrite will no longer run for tf.functions, and so float32 operations will no longer be converted to float16. This does not undo the effects of loss scaling. Any optimizers wrapped with a LossScaleOptimizer will continue to do loss scaling, although this loss scaling will no longer be useful, as the graph rewrite no longer converts tf.functions to use float16. This function is useful for unit testing. A unit test can test using the mixed precision graph rewrite, then disable it so future unit tests continue using float32. """ if not mixed_precision_global_state.mixed_precision_graph_rewrite_is_enabled: tf_logging.warn('disable_mixed_precision_graph_rewrite() called when mixed ' 'precision is already disabled.') config.set_optimizer_experimental_options({'auto_mixed_precision': False}) mixed_precision_global_state.mixed_precision_graph_rewrite_is_enabled = False
def disable_mixed_precision_graph_rewrite(): """Disables the mixed precision graph rewrite. After this is called, the mixed precision graph rewrite will no longer run for tf.functions, and so float32 operations will no longer be converted to float16. This does not undo the effects of loss scaling. Any optimizers wrapped with a LossScaleOptimizer will continue to do loss scaling, although this loss scaling will no longer be useful, as the graph rewrite no longer converts tf.functions to use float16. This function is useful for unit testing. A unit test can test using the mixed precision graph rewrite, then disable it so future unit tests continue using float32. """ if not mixed_precision_global_state.mixed_precision_is_enabled: tf_logging.warn('disable_mixed_precision_graph_rewrite() called when mixed ' 'precision is already disabled.') config.set_optimizer_experimental_options({'auto_mixed_precision': False}) mixed_precision_global_state.mixed_precision_is_enabled = False
def testOptimizerToggleOptionPinToHost(self): options = config.get_optimizer_experimental_options() self.assertIsNone(options.get('pin_to_host_optimization')) @def_function.function def fun(): op = test_ops.device_placement_op() return op # Force optimizer to run for all graphs config.set_optimizer_experimental_options({'min_graph_nodes': -1}) options['min_graph_nodes'] = -1 # Since pin to host is disabled, the operation should go on GPU gpu = self.evaluate(fun()) self.assertIn(compat.as_bytes('GPU'), gpu) config.set_optimizer_experimental_options( {'pin_to_host_optimization': True}) options['pin_to_host_optimization'] = True self.assertDictEqual(config.get_optimizer_experimental_options(), options) self.assertDictEqual( context.context().get_optimizer_experimental_options(), options) # Since pin to host is enabled, the operation should go on CPU cpu = self.evaluate(fun()) self.assertIn(compat.as_bytes('CPU'), cpu) config.set_optimizer_experimental_options( {'pin_to_host_optimization': False}) options['pin_to_host_optimization'] = False self.assertDictEqual(config.get_optimizer_experimental_options(), options) self.assertDictEqual( context.context().get_optimizer_experimental_options(), options) # Since pin to host is disabled again, the operation should go on GPU gpu2 = self.evaluate(fun()) self.assertIn(compat.as_bytes('GPU'), gpu2)
def testOptimizerToggleOptionPinToHost(self): options = config.get_optimizer_experimental_options() self.assertIsNone(options.get('pin_to_host_optimization')) @def_function.function def fun(): op = test_ops.device_placement_op() return op # Force optimizer to run for all graphs config.set_optimizer_experimental_options({'min_graph_nodes': -1}) options['min_graph_nodes'] = -1 # Since pin to host is disabled, the operation should go on GPU gpu = self.evaluate(fun()) self.assertIn(compat.as_bytes('GPU'), gpu) config.set_optimizer_experimental_options( {'pin_to_host_optimization': True}) options['pin_to_host_optimization'] = True self.assertDictEqual(config.get_optimizer_experimental_options(), options) self.assertDictEqual( context.context().get_optimizer_experimental_options(), options) # Since pin to host is enabled, the operation should go on CPU cpu = self.evaluate(fun()) self.assertIn(compat.as_bytes('CPU'), cpu) config.set_optimizer_experimental_options( {'pin_to_host_optimization': False}) options['pin_to_host_optimization'] = False self.assertDictEqual(config.get_optimizer_experimental_options(), options) self.assertDictEqual( context.context().get_optimizer_experimental_options(), options) # Since pin to host is disabled again, the operation should go on GPU gpu2 = self.evaluate(fun()) self.assertIn(compat.as_bytes('GPU'), gpu2)
def enable_mixed_precision_graph_rewrite_v1(opt, loss_scale='dynamic'): """Enable mixed precision via a graph rewrite. Mixed precision is the use of both float32 and float16 data types when training a model to improve performance. This is achieved via a graph rewrite operation and a loss-scale optimizer. Performing arithmetic operations in float16 takes advantage of specialized processing units, such as NVIDIA Tensor Cores, for much higher arithmetic throughput. However, due to the smaller representable range, performing the entire training with float16 can result in gradient underflow, that is, small gradient values becoming zeroes. Instead, performing only select arithmetic operations in float16 results in higher throughput and decreased training time when using compatible hardware accelerators while also reducing memory usage, typically without sacrificing model accuracy. Note: While the mixed precision rewrite changes the datatype of various layers throughout the model, the same accuracy reached in float32 is expected. If a `NaN` gradient occurs with dynamic loss scaling, the model update for that batch is skipped. In this case, the global step count is not incremented, and the `LossScaleOptimizer` attempts to decrease the loss scaling value to avoid `NaN` values in subsequent iterations. This approach has been shown to achieve the same accuracy as float32 and, in most cases, better training throughput. Example: ```python model = tf.keras.models.Sequential([ tf.keras.layers.Dense(64, activation='relu'), tf.keras.layers.Dense(64, activation='softmax'), ]) opt = tf.keras.optimizers.SGD() opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt) model.compile(loss="mse", optimizer=opt) x_train = np.random.random((1024, 64)) y_train = np.random.random((1024, 64)) model.fit(x_train, y_train) ``` Calling `enable_mixed_precision_graph_rewrite(opt)` enables the graph rewrite operation before computing gradients. The function additionally returns an `Optimizer` (`opt`) wrapped with a `LossScaleOptimizer`. This prevents underflow in the float16 tensors during the backward pass. An optimizer of type `tf.train.Optimizer` or `tf.keras.optimizers.Optimizer` must be passed to this function, which will then be wrapped to use loss scaling. The graph rewrite operation changes the `dtype` of certain operations in the graph from float32 to float16. There are several categories of operations that are either included or excluded by this rewrite operation. The following categories of Ops are defined inside corresponding functions under the class `AutoMixedPrecisionLists` in <a href="https://github.com/tensorflow/tensorflow/blob/master/tensorflow/ core/grappler/optimizers/auto_mixed_precision_lists.h"> auto_mixed_precision_lists.h</a>: * `ClearList`: Ops that do not have numerically significant adverse effects. E.g. `ArgMax` and `Floor`. * `AllowList`: Ops that are considered numerically safe for execution in float16, and thus are always converted. E.g. `Conv2D`. * `DenyList`: Ops that are numerically unsafe to execute in float16 and can negatively affect downstream nodes. E.g. `Softmax`. * `GrayList`: Ops that are considered numerically safe for execution in float16 unless downstream from a DenyList Op. E.g. `Add` and `AvgPool`. When this function is used, gradients should only be computed and applied with the returned optimizer, either by calling `opt.minimize()` or `opt.compute_gradients()` followed by `opt.apply_gradients()`. Gradients should not be computed with `tf.gradients` or `tf.GradientTape`. This is because the returned optimizer will apply loss scaling, and `tf.gradients` or `tf.GradientTape` will not. If you do directly use `tf.gradients` or `tf.GradientTape`, your model may not converge due to float16 underflow problems. When eager execution is enabled, the mixed precision graph rewrite is only enabled within `tf.function`s, as outside `tf.function`s, there is no graph. For NVIDIA GPUs with Tensor cores, as a general performance guide, dimensions (such as batch size, input size, output size, and channel counts) should be powers of two if under 256, or otherwise divisible by 8 if above 256. For more information, check out the [NVIDIA Deep Learning Performance Guide]( https://docs.nvidia.com/deeplearning/sdk/dl-performance-guide/index.html). Currently, mixed precision is only enabled on NVIDIA Tensor Core GPUs with Compute Capability 7.0 and above (Volta, Turing, or newer architectures). The parts of the graph on CPUs and TPUs are untouched by the graph rewrite. Raises: `ValueError`, if the `tf.keras.mixed_precision` API is also used by calling `tf.keras.mixed_precision.experimental.set_policy`. Only one mixed precision API can be used. Args: opt: An instance of a `tf.keras.optimizers.Optimizer` or a `tf.train.Optimizer`. loss_scale: Either an int/float, the string `"dynamic"`, or an instance of a `tf.mixed_precision.experimental.LossScale`. The loss scale to use. It is recommended to keep this as its default value of `"dynamic"`, which will adjust the scaling automatically to prevent `Inf` or `NaN` values. Returns: A version of `opt` that will use loss scaling to prevent underflow. """ if mixed_precision_global_state.is_using_mixed_precision_policy(): raise ValueError( 'The mixed precision graph rewrite cannot be enabled, because the ' 'global Keras dtype Policy has been set to a mixed precision policy. ' 'At most, one of the following can be called:\n\n' ' 1. tf.keras.mixed_precision.experimental.set_policy() with a mixed ' 'precision policy (You called this first)\n\n' ' 2. tf.train.experimental.enable_mixed_precision_graph_rewrite() ' '(You called this second)\n' 'You called both functions, which is an error, because both functions ' 'enable you to use mixed precision. If in doubt which function to use, ' 'use the first, as it supports Eager execution and is more ' 'customizable.') if mixed_precision_global_state.non_mixed_precision_session_created(): # TODO(reedwm): Give the stacktrace of the existing Sessions. And if the # Sessions have already been closed, do not raise this error message. tf_logging.warn('You already have existing Sessions that do not use mixed ' 'precision. enable_mixed_precision_graph_rewrite() will ' 'not affect these Sessions.') opt = _wrap_optimizer(opt, loss_scale) config.set_optimizer_experimental_options({'auto_mixed_precision': True}) mixed_precision_global_state.set_mixed_precision_graph_rewrite_enabled(True) return opt
def _enable_mixed_precision_graph_rewrite_base(opt, loss_scale, use_v1_behavior): """Enables mixed precision. See `enable_mixed_precision_graph_rewrite`.""" opt = _wrap_optimizer(opt, loss_scale, use_v1_behavior=use_v1_behavior) config.set_optimizer_experimental_options({'auto_mixed_precision': True}) return opt