wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/python/keras/mixed_precision/BUILD
+153
View File
@@ -0,0 +1,153 @@
# Copyright 2018 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.
# ==============================================================================
# Description:
# Contains the Keras Mixed Precision API (TensorFlow version).
load("@xla//third_party/rules_python/python:defs.bzl", "py_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
# TODO(scottzhu): Remove these two deps and convert the test to integration test.
"//tensorflow/python/distribute:__pkg__", # For collective_all_reduce_strategy_test
"//tensorflow/python/keras:__subpackages__",
"//tensorflow/tools/pip_package:__pkg__",
"//tensorflow/tools/pip_package:__subpackages__",
],
licenses = ["notice"],
)
filegroup(
name = "all_py_srcs",
srcs = glob(["*.py"]),
visibility = ["//tensorflow/python/keras/google/private_tf_api_test:__pkg__"],
)
py_library(
name = "mixed_precision_experimental",
srcs = ["__init__.py"],
srcs_version = "PY3",
strict_deps = False,
deps = [
":get_layer_policy",
":loss_scale_optimizer",
":policy",
],
)
py_library(
name = "policy",
srcs = [
"policy.py",
],
srcs_version = "PY3",
strict_deps = False,
deps = [
":device_compatibility_check",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/training/experimental:mixed_precision_global_state",
],
)
py_library(
name = "device_compatibility_check",
srcs = ["device_compatibility_check.py"],
srcs_version = "PY3",
strict_deps = False,
deps = [
"//tensorflow/python/framework:config",
],
)
py_library(
name = "get_layer_policy",
srcs = ["get_layer_policy.py"],
srcs_version = "PY3",
strict_deps = False,
deps = [
"//tensorflow/python/keras/engine:base_layer",
],
)
py_library(
name = "autocast_variable",
srcs = [
"autocast_variable.py",
],
srcs_version = "PY3",
strict_deps = False,
deps = [
"//tensorflow/python/distribute:distribute_utils",
"//tensorflow/python/distribute:ps_values",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_conversion",
"//tensorflow/python/framework:tensor_conversion_registry",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/keras/distribute",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/types:core",
],
)
py_library(
name = "loss_scale",
srcs = ["loss_scale.py"],
srcs_version = "PY3",
strict_deps = False,
deps = [
"//tensorflow/python/keras/utils:generic_utils",
"//tensorflow/python/training/experimental:loss_scale",
],
)
py_library(
name = "loss_scale_optimizer",
srcs = ["loss_scale_optimizer.py"],
srcs_version = "PY3",
strict_deps = False,
deps = [
":loss_scale",
"//tensorflow/python/distribute:collective_all_reduce_strategy",
"//tensorflow/python/distribute:distribute_lib",
"//tensorflow/python/distribute:mirrored_strategy",
"//tensorflow/python/distribute:one_device_strategy",
"//tensorflow/python/distribute:tpu_strategy",
"//tensorflow/python/framework:tensor_conversion",
"//tensorflow/python/keras/optimizer_v2",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:variable_v1",
"//tensorflow/python/trackable:base_delegate",
"@absl_py//absl/testing:parameterized",
],
)
py_library(
name = "test_util",
srcs = ["test_util.py"],
srcs_version = "PY3",
strict_deps = False,
deps = [
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_conversion",
"//tensorflow/python/keras",
"//tensorflow/python/ops:cond",
],
)
tensorflow/python/keras/mixed_precision/__init__.py
+20
View File
@@ -0,0 +1,20 @@
# 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.
# ==============================================================================
"""Keras mixed precision API.
See [the mixed precision guide](
https://www.tensorflow.org/guide/keras/mixed_precision) to learn how to
use the API.
"""
tensorflow/python/keras/mixed_precision/autocast_variable.py
+560
View File
@@ -0,0 +1,560 @@
# Copyright 2019 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.
# ==============================================================================
"""Contains AutoCastVariable, a variable which automatically casts itself."""
import threading
from tensorflow.python.eager import context
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_conversion
from tensorflow.python.framework import tensor_conversion_registry
from tensorflow.python.framework import tensor_shape
from tensorflow.python.keras.distribute import distributed_training_utils
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import variables
from tensorflow.python.types import core
# _autocast_dtype.dtype is the dtype AutoCastVariables should be cast to, or
# None if AutoCastVariables should not be cast.
_autocast_dtype = threading.local()
def numpy_text(tensor, is_repr=False):
"""Human readable representation of a tensor's numpy value."""
if tensor.dtype.is_numpy_compatible:
# pylint: disable=protected-access
text = repr(tensor._numpy()) if is_repr else str(tensor._numpy())
# pylint: enable=protected-access
else:
text = '<unprintable>'
if '\n' in text:
text = '\n' + text
return text
class AutoCastVariable(variables.Variable, core.Tensor):
"""Variable that will cast itself to a different dtype in applicable contexts.
This class wraps a floating-point `tf.Variable`. It emulates the variable
interface and delegates to the wrapped variable, but it additionally will cast
the wrapped variable under an `enable_auto_cast_variables(dtype)` context
manager.
For example:
>>> v = tf.Variable(1.0, dtype=tf.float32)
>>> v = AutoCastVariable(v)
>>> tf.identity(v).dtype
tf.float32
>>> with enable_auto_cast_variables(tf.float16):
... tf.identity(v).dtype
tf.float16
The purpose of this class is to allow Keras layers to create variables in
float32, and automatically cast them to float16 or bfloat16 when the layer is
called.
"""
def __init__(self, variable):
"""Creates an AutoCastVariable instance.
Args:
variable: A floating-point resource variable to wrap.
Raises:
ValueError: If `variable` is not a floating-point resource variable
"""
if not isinstance(variable, variables.Variable):
raise ValueError('variable must be of type tf.ResourceVariable, but got: '
'%s' % variable)
if not variable.dtype.is_floating:
raise ValueError('variable must be a floating point variable but has '
'type: %s' % variable.dtype.name)
self._variable = variable
# 'delegate' means AutoCastVariable.op return self._variable.op, which will
# raise an AttributeError in Eager (as intended). If set to any other value,
# AutoCastVariable.op returns that value instead, which is used to set the
# op attribute in AutoCastVariable.assign().
self._op = 'delegate'
def _should_cast(self):
"""Returns True if this variable should be casted when accessed."""
autocast_dtype = getattr(_autocast_dtype, 'dtype', None)
return autocast_dtype is not None and self.dtype != autocast_dtype
@property
def dtype(self):
"""The dtype of the underlying variable, before any casts are done."""
return self._variable.dtype
@property
def true_dtype(self):
"""Deprecated alias of `dtype`."""
return self._variable.dtype
@property
def _cast_dtype(self):
dtype = getattr(_autocast_dtype, 'dtype', None)
return dtype or self._variable.dtype
def value(self):
val = self._variable.value()
if not self._should_cast():
return val
return math_ops.cast(val, self._cast_dtype)
def read_value(self):
val = self._variable.read_value()
return math_ops.cast(val, self._cast_dtype)
def sparse_read(self, indices, name=None):
"""Reads the value of this variable sparsely, using `gather`."""
val = self._variable.sparse_read(indices, name=name)
return math_ops.cast(val, self._cast_dtype)
def gather_nd(self, indices, name=None):
"""Gather slices of the variable into a Tensor."""
val = self._variable.gather_nd(indices, name=name)
return math_ops.cast(val, self._cast_dtype)
def __getattr__(self, name):
return getattr(self._variable, name)
def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False):
"""Converts this variable to a tensor."""
if as_ref:
# This ValueError should not occur in practice since it is impossible to
# pass as_ref=True using public APIs.
raise ValueError('Cannot convert AutoCastVariable to a tensor if '
'as_ref=True is passed to convert_to_tensor')
if not self._should_cast():
return tensor_conversion.convert_to_tensor_v2_with_dispatch(
self._variable, dtype=dtype, name=name
)
if dtype is not None and not dtype.is_compatible_with(self._cast_dtype):
raise ValueError(
'Incompatible type conversion requested to type {!r} for '
'AutoCastVariable which is casted to type {!r}'.format(
dtype.name, self._cast_dtype.name))
val = tensor_conversion.convert_to_tensor_v2_with_dispatch(
self._variable, dtype=self._variable.dtype, name=name
)
return math_ops.cast(val, self._cast_dtype)
def _should_act_as_resource_variable(self):
"""Pass resource_variable_ops.is_resource_variable check."""
pass
def __repr__(self):
if context.executing_eagerly() and not self._in_graph_mode:
repr_str = ("<AutoCastVariable '{v.name}' shape={v.shape} "
'dtype={v.dtype.name} dtype_to_cast_to={v._cast_dtype.name}, '
'numpy={np_repr}>')
return repr_str.format(
v=self, np_repr=numpy_text(self.read_value(), is_repr=True))
else:
repr_str = ("<AutoCastVariable '{v.name}' shape={v.shape} "
'dtype={v.dtype.name} dtype_to_cast_to={v._cast_dtype.name}>')
return repr_str.format(v=self)
# Method delegations: We delegate the following methods to self._variable.
# Each of these methods simply calls the same method on self._variable. The
# base Variable raises NotImplementedError for most of these, so we must
# override them.
#
# We do not define the following methods from Variable for the following
# reasons:
# * 'count_up_to': This method only applies to int variables, which cannot
# be wrapped with an AutoCastVariable.
# * 'ref': Instead we inherit the definition from Variable.
# If we defined and delegated to Variable, the ref of an AutoCastVariable
# would be the same as the ref of the underlying variable, which would be
# strange as they are different Python objects.
def set_shape(self, shape):
return self._variable.set_shape(self, shape)
@property
def trainable(self):
return self._variable.trainable
@property
def synchronization(self):
return self._variable.synchronization
@property
def aggregation(self):
return self._variable.aggregation
def eval(self, session=None):
return self._variable.eval(session)
def initialized_value(self):
return self._variable.initialized_value()
@property
def initial_value(self):
return self._variable.initial_value
@property
def constraint(self):
return self._variable.constraint
def _apply_assign_update(self,
update_fn,
value,
use_locking=None,
name=None,
read_value=True):
# TODO(b/146181571): This logic can be simplified once
# DistributedVariable.assign returns a DistributedVariable. Currently for
# MirroredStrategy, it returns a Mirrored value.
if ops.executing_eagerly_outside_functions():
assign_op = update_fn(value, use_locking, name, False)
if read_value:
# We create a new AutoCastVariable with the same underlying tf.Variable.
# The new AutoCastVariable is identical except the 'op' attribute is
# defined. This matches the behavior of tf.Variable.assign.
var = create_autocast_variable(self._variable)
var._op = assign_op # pylint:disable=protected-access
return var
return assign_op
# Fallback to wrapping the returned variable in graph mode if possible
assign_var = update_fn(value, use_locking, name, read_value)
if read_value and resource_variable_ops.is_resource_variable(assign_var):
return create_autocast_variable(assign_var)
return assign_var
def _apply_update(self, update_fn, *args, **kwargs):
update_var = update_fn(*args, **kwargs)
if ops.executing_eagerly_outside_functions():
return self
# Fallback to wrapping the returned variable in graph mode if possible
if resource_variable_ops.is_resource_variable(update_var):
return create_autocast_variable(update_var)
return update_var
def assign(self, value, use_locking=None, name=None, read_value=True):
return self._apply_assign_update(self._variable.assign, value, use_locking,
name, read_value)
def assign_add(self, delta, use_locking=None, name=None, read_value=True):
return self._apply_assign_update(self._variable.assign_add, delta,
use_locking, name, read_value)
def assign_sub(self, delta, use_locking=None, name=None, read_value=True):
return self._apply_assign_update(self._variable.assign_sub, delta,
use_locking, name, read_value)
def scatter_sub(self, sparse_delta, use_locking=False, name=None):
return self._apply_update(self._variable.scatter_sub, sparse_delta,
use_locking, name)
def scatter_add(self, sparse_delta, use_locking=False, name=None):
return self._apply_update(self._variable.scatter_add, sparse_delta,
use_locking, name)
def scatter_max(self, sparse_delta, use_locking=False, name=None):
return self._apply_update(self._variable.scatter_max, sparse_delta,
use_locking, name)
def scatter_min(self, sparse_delta, use_locking=False, name=None):
return self._apply_update(self._variable.scatter_min, sparse_delta,
use_locking, name)
def scatter_mul(self, sparse_delta, use_locking=False, name=None):
return self._apply_update(self._variable.scatter_mul, sparse_delta,
use_locking, name)
def scatter_div(self, sparse_delta, use_locking=False, name=None):
return self._apply_update(self._variable.scatter_div, sparse_delta,
use_locking, name)
def scatter_update(self, sparse_delta, use_locking=False, name=None):
return self._apply_update(self._variable.scatter_update, sparse_delta,
use_locking, name)
def batch_scatter_update(self, sparse_delta, use_locking=False, name=None):
return self._apply_update(self._variable.batch_scatter_update, sparse_delta,
use_locking, name)
def scatter_nd_sub(self, indices, updates, name=None):
return self._apply_update(self._variable.scatter_nd_sub, indices, updates,
name)
def scatter_nd_add(self, indices, updates, name=None):
return self._apply_update(self._variable.scatter_nd_add, indices, updates,
name)
def scatter_nd_update(self, indices, updates, name=None):
return self._apply_update(self._variable.scatter_nd_update, indices,
updates, name)
def load(self, value, session=None):
return self._variable.load(value, session)
@property
def name(self):
return self._variable.name
@property
def _shared_name(self):
return self._variable._shared_name # pylint:disable=protected-access
@property
def initializer(self):
return self._variable.initializer
@property
def device(self):
return self._variable.device
@property
def op(self):
if self._op == 'delegate':
return self._variable.op
return self._op
def _as_graph_element(self):
graph_element = self._variable._as_graph_element() # pylint:disable=protected-access
if graph_element is None:
return self._op
return graph_element
@property
def graph(self):
return self._variable.graph
@property
def shape(self):
return self._variable.shape
def get_shape(self) -> tensor_shape.TensorShape:
return self._variable.get_shape()
def _gather_saveables_for_checkpoint(self):
# By delegating this method to the wrapped variable, checkpoints with
# AutoCastVariables are identical to checkpoints with normal variables.
# Therefore models checkpointed with AutoCastVariables can be restored on
# models with normal variables, and vice versa.
return self._variable._gather_saveables_for_checkpoint() # pylint:disable=protected-access
def _export_to_saved_model_graph(self, object_map, tensor_map, options,
**kwargs):
# By delegating this method to the wrapped variable, SavedModel with
# AutoCastVariables are identical to SavedModel with normal variables.
resource_list = self._variable._export_to_saved_model_graph( # pylint:disable=protected-access
object_map, tensor_map, options, **kwargs)
object_map[self] = object_map[self._variable]
return resource_list
# TODO(reedwm): Maybe encode the fact the variable is an AutoCastVariable in
# to_proto().
def to_proto(self, export_scope=None):
return self._variable.to_proto(export_scope)
def from_proto(self, variable_def, import_scope=None):
return self._variable.from_proto(variable_def, import_scope)
# Delegate the private attributes _handle_name and _initializer_op to
# self._variable. SavedModel sets these attributes when loading a model. For
# example, it sets _handle_name here:
# https://github.com/tensorflow/tensorflow/blob/db26bd574fa95b5bdd53c08463dd19407cc0297e/tensorflow/python/keras/saving/saved_model/load.py#L211
# We need to expose these attributes on AutoCastVariable as well for
# SavedModel to work properly.
# TODO(reedwm/kathywu): Find a better way to support SavedModel. Exposing
# private attributes is hacky and difficult to maintain.
@property
def _handle_name(self):
return self._variable._handle_name # pylint: disable=protected-access
@_handle_name.setter
def _handle_name(self, handle_name):
self._variable._handle_name = handle_name # pylint: disable=protected-access
@property
def _initializer_op(self):
return self._variable._initializer_op # pylint: disable=protected-access
@_initializer_op.setter
def _initializer_op(self, initializer_op):
self._variable._initializer_op = initializer_op # pylint: disable=protected-access
# Operator overloads:
# Note we only overload operators that support floating-point types, as
# non-float variables cannot be wrapped with an AutoCastVariable.
# Also note: We call read_value() instead of value(), because value() causes
# gradients not to work properly when TPUStrategy is used: b/143380936
def __add__(self, o):
return self.read_value() + o
def __radd__(self, o):
return o + self.read_value()
def __sub__(self, o):
return self.read_value() - o
def __rsub__(self, o):
return o - self.read_value()
def __mul__(self, o):
return self.read_value() * o
def __rmul__(self, o):
return o * self.read_value()
def __truediv__(self, o):
return self.read_value() / o
def __rtruediv__(self, o):
return o / self.read_value()
def __floordiv__(self, o):
return self.read_value() // o
def __rfloordiv__(self, o):
return o // self.read_value()
def __mod__(self, o):
return self.read_value() % o
def __rmod__(self, o):
return o % self.read_value()
def __lt__(self, o):
return self.read_value() < o
def __le__(self, o):
return self.read_value() <= o
def __gt__(self, o):
return self.read_value() > o
def __ge__(self, o):
return self.read_value() >= o
def __getitem__(self, o):
return self.read_value()[o]
def __pow__(self, o, modulo=None):
return pow(self.read_value(), o, modulo)
def __rpow__(self, o):
return pow(o, self.read_value())
def __neg__(self):
return -self.read_value() # pylint: disable=invalid-unary-operand-type
def __abs__(self):
return abs(self.read_value())
def __div__(self, o):
try:
return self.read_value().__div__(o)
except AttributeError:
# See https://docs.python.org/3/library/constants.html#NotImplemented
return NotImplemented
def __rdiv__(self, o):
try:
return self.read_value().__rdiv__(o)
except AttributeError:
# See https://docs.python.org/3/library/constants.html#NotImplemented
return NotImplemented
def __matmul__(self, o):
try:
return self.read_value().__matmul__(o)
except AttributeError:
# See https://docs.python.org/3/library/constants.html#NotImplemented
return NotImplemented
def __rmatmul__(self, o):
try:
return self.read_value().__rmatmul__(o)
except AttributeError:
# See https://docs.python.org/3/library/constants.html#NotImplemented
return NotImplemented
# pylint: enable=multiple-statements
tensor_conversion_registry.register_tensor_conversion_function(
AutoCastVariable, AutoCastVariable._dense_var_to_tensor) # pylint:disable=protected-access
def create_autocast_variable(variable):
"""Creates an AutoCastVariable that wraps another variable.
This typically just returns `AutoCastVariable(variable)`. But, if the variable
is a DistributedVariable or one of its subclasses, we instead dynamically
create a class that subclasses from both AutoCastVariable and
variable.__class__. This is so the returned variable will still pass
`isinstance(variable, variable.__class__)`, which is required for
DistributedVariables and its subclasses to work properly.
Args:
variable: A floating-point resource variable to wrap.
Returns:
An AutoCastVariable that wraps the variable.
"""
if not distributed_training_utils.is_distributed_variable(variable):
return AutoCastVariable(variable)
class AutoCastDistributedVariable(AutoCastVariable, variable.__class__):
"""An AutoCastVariable that also subclasses from variable.__class__.
variable.__class__ is either a DistributedVariable or an
AggregatingVariable.
"""
def __repr__(self):
# pylint: disable=missing-format-attribute
return ('<AutoCastDistributedVariable dtype={v.dtype.name} '
'dtype_to_cast_to={v._cast_dtype.name} '
'inner_variable={v._variable}>'
).format(v=self)
# pylint: enable=missing-format-attribute
return AutoCastDistributedVariable(variable)
class enable_auto_cast_variables(object): # pylint:disable=invalid-name
"""Context manager which enables the autocasting of `AutoCastVariable`s.
Under this context manager, `AutoCastVariable`s will be cast to `dtype` if
`dtype` is floating-point. Otherwise, `AutoCastVariable`s will not be cast.
"""
__slots__ = ['_dtype', '_prev_dtype']
def __init__(self, dtype):
if dtype and not dtype.is_floating:
dtype = None
self._dtype = dtype
def __enter__(self):
self._prev_dtype = getattr(_autocast_dtype, 'dtype', None)
_autocast_dtype.dtype = self._dtype
def __exit__(self, type_arg, value_arg, traceback_arg):
_autocast_dtype.dtype = self._prev_dtype
tensorflow/python/keras/mixed_precision/device_compatibility_check.py
+147
View File
@@ -0,0 +1,147 @@
# Copyright 2019 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.
# ==============================================================================
"""Contains function to log if devices are compatible with mixed precision."""
import itertools
from tensorflow.python.framework import config
from tensorflow.python.platform import tf_logging
_COMPAT_CHECK_PREFIX = 'Mixed precision compatibility check (mixed_float16): '
_COMPAT_CHECK_OK_PREFIX = _COMPAT_CHECK_PREFIX + 'OK'
_COMPAT_CHECK_WARNING_PREFIX = _COMPAT_CHECK_PREFIX + 'WARNING'
_COMPAT_CHECK_WARNING_SUFFIX = (
'If you will use compatible GPU(s) not attached to this host, e.g. by '
'running a multi-worker model, you can ignore this warning. This message '
'will only be logged once')
def _dedup_strings(device_strs):
"""Groups together consecutive identical strings.
For example, given:
['GPU 1', 'GPU 2', 'GPU 2', 'GPU 3', 'GPU 3', 'GPU 3']
This function returns:
['GPU 1', 'GPU 2 (x2)', 'GPU 3 (x3)']
Args:
device_strs: A list of strings, each representing a device.
Returns:
A copy of the input, but identical consecutive strings are merged into a
single string.
"""
new_device_strs = []
for device_str, vals in itertools.groupby(device_strs):
num = len(list(vals))
if num == 1:
new_device_strs.append(device_str)
else:
new_device_strs.append('%s (x%d)' % (device_str, num))
return new_device_strs
def _log_device_compatibility_check(policy_name, gpu_details_list):
"""Logs a compatibility check if the devices support the policy.
Currently only logs for the policy mixed_float16.
Args:
policy_name: The name of the dtype policy.
gpu_details_list: A list of dicts, one dict per GPU. Each dict
is the device details for a GPU, as returned by
`tf.config.experimental.get_device_details()`.
"""
if policy_name != 'mixed_float16':
# TODO(b/145686977): Log if the policy is 'mixed_bfloat16'. This requires
# checking if a TPU is available.
return
supported_device_strs = []
unsupported_device_strs = []
for details in gpu_details_list:
name = details.get('device_name', 'Unknown GPU')
cc = details.get('compute_capability')
if cc:
device_str = '%s, compute capability %s.%s' % (name, cc[0], cc[1])
if cc >= (7, 0):
supported_device_strs.append(device_str)
else:
unsupported_device_strs.append(device_str)
else:
unsupported_device_strs.append(
name + ', no compute capability (probably not an Nvidia GPU)')
if unsupported_device_strs:
warning_str = _COMPAT_CHECK_WARNING_PREFIX + '\n'
if supported_device_strs:
warning_str += ('Some of your GPUs may run slowly with dtype policy '
'mixed_float16 because they do not all have compute '
'capability of at least 7.0. Your GPUs:\n')
elif len(unsupported_device_strs) == 1:
warning_str += ('Your GPU may run slowly with dtype policy mixed_float16 '
'because it does not have compute capability of at least '
'7.0. Your GPU:\n')
else:
warning_str += ('Your GPUs may run slowly with dtype policy '
'mixed_float16 because they do not have compute '
'capability of at least 7.0. Your GPUs:\n')
for device_str in _dedup_strings(supported_device_strs +
unsupported_device_strs):
warning_str += ' ' + device_str + '\n'
warning_str += ('See https://developer.nvidia.com/cuda-gpus for a list of '
'GPUs and their compute capabilities.\n')
warning_str += _COMPAT_CHECK_WARNING_SUFFIX
tf_logging.warning(warning_str)
elif not supported_device_strs:
tf_logging.warning(
'%s\n'
'The dtype policy mixed_float16 may run slowly because '
'this machine does not have a GPU. Only Nvidia GPUs with '
'compute capability of at least 7.0 run quickly with '
'mixed_float16.\n%s' % (_COMPAT_CHECK_WARNING_PREFIX,
_COMPAT_CHECK_WARNING_SUFFIX))
elif len(supported_device_strs) == 1:
tf_logging.info('%s\n'
'Your GPU will likely run quickly with dtype policy '
'mixed_float16 as it has compute capability of at least '
'7.0. Your GPU: %s' % (_COMPAT_CHECK_OK_PREFIX,
supported_device_strs[0]))
else:
tf_logging.info('%s\n'
'Your GPUs will likely run quickly with dtype policy '
'mixed_float16 as they all have compute capability of at '
'least 7.0' % _COMPAT_CHECK_OK_PREFIX)
_logged_compatibility_check = False
def log_device_compatibility_check(policy_name):
"""Logs a compatibility check if the devices support the policy.
Currently only logs for the policy mixed_float16. A log is shown only the
first time this function is called.
Args:
policy_name: The name of the dtype policy.
"""
global _logged_compatibility_check
if _logged_compatibility_check:
return
_logged_compatibility_check = True
gpus = config.list_physical_devices('GPU')
gpu_details_list = [config.get_device_details(g) for g in gpus]
_log_device_compatibility_check(policy_name, gpu_details_list)
tensorflow/python/keras/mixed_precision/get_layer_policy.py
+39
View File
@@ -0,0 +1,39 @@
# Copyright 2019 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.
# ==============================================================================
"""Contains the get_layer_policy function.
This is a separate file from policy.py to avoid a circular dependency.
get_layer_policy() relies on base_layer.py, itself which relies on policy.py.
"""
from tensorflow.python.keras.engine import base_layer
def get_layer_policy(layer):
"""Returns the dtype policy of a layer.
Warning: This function is deprecated. Use
`tf.keras.layers.Layer.dtype_policy` instead.
Args:
layer: A `tf.keras.layers.Layer`.
Returns:
The `tf.keras.mixed_precision.Policy` of the layer.
"""
if not isinstance(layer, base_layer.Layer):
raise ValueError('get_policy can only be called on a layer, but got: %s'
% (layer,))
return layer.dtype_policy
tensorflow/python/keras/mixed_precision/loss_scale.py
+58
View File
@@ -0,0 +1,58 @@
# Copyright 2019 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.
# ==============================================================================
"""Contains keras-specific LossScale functionality.
This functions cannot be in the non-keras loss_scale.py file since they depend
on keras, and files outside of keras should not depend on files inside keras.
"""
from tensorflow.python.keras.utils import generic_utils
from tensorflow.python.training.experimental import loss_scale as loss_scale_module
def serialize(loss_scale):
return generic_utils.serialize_keras_object(loss_scale)
def deserialize(config, custom_objects=None):
loss_scale_module_objects = {
'FixedLossScale': loss_scale_module.FixedLossScale,
'DynamicLossScale': loss_scale_module.DynamicLossScale,
}
return generic_utils.deserialize_keras_object(
config,
module_objects=loss_scale_module_objects,
custom_objects=custom_objects,
printable_module_name='loss scale'
)
def get(identifier):
"""Get a loss scale object."""
if isinstance(identifier, dict):
return deserialize(identifier)
if isinstance(identifier, (int, float)):
return loss_scale_module.FixedLossScale(identifier)
if identifier == 'dynamic':
return loss_scale_module.DynamicLossScale()
if isinstance(identifier, loss_scale_module.LossScale):
return identifier
elif identifier is None:
return None
else:
raise ValueError('Could not interpret loss scale identifier: %s' %
identifier)
tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py
+1136
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/keras/mixed_precision/policy.py
+582
View File
@@ -0,0 +1,582 @@
# Copyright 2019 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.
# ==============================================================================
"""Contains the Policy class for mixed precision training."""
import contextlib
from tensorflow.python.framework import dtypes
from tensorflow.python.keras import backend
from tensorflow.python.keras.engine import base_layer_utils
from tensorflow.python.keras.mixed_precision import device_compatibility_check
from tensorflow.python.keras.mixed_precision import loss_scale as keras_loss_scale_module
from tensorflow.python.keras.utils import generic_utils
from tensorflow.python.platform import tf_logging
from tensorflow.python.training.experimental import mixed_precision_global_state
# pylint: disable=g-classes-have-attributes
class Policy(object):
"""A dtype policy for a Keras layer.
A dtype policy determines a layer's computation and variable dtypes. Each
layer has a policy. Policies can be passed to the `dtype` argument of layer
constructors, or a global policy can be set with
`tf.keras.mixed_precision.set_global_policy`.
Args:
name: The policy name, which determines the compute and variable dtypes. Can
be any dtype name, such as `'float32'` or `'float64'`, which causes both
the compute and variable dtypes will be that dtype. Can also be the string
`'mixed_float16'` or `'mixed_bfloat16'`, which causes the compute dtype to
be float16 or bfloat16 and the variable dtype to be float32.
Typically you only need to interact with dtype policies when using mixed
precision, which is the use of float16 or bfloat16 for computations and
float32 for variables. This is why the term `mixed_precision` appears in the
API name. Mixed precision can be enabled by passing `'mixed_float16'` or
`'mixed_bfloat16'` to `tf.keras.mixed_precision.set_global_policy`. See [the
mixed precision guide](https://www.tensorflow.org/guide/keras/mixed_precision)
for more information on how to use mixed precision.
>>> tf.keras.mixed_precision.set_global_policy('mixed_float16')
>>> layer1 = tf.keras.layers.Dense(10)
>>> layer1.dtype_policy # `layer1` will automatically use mixed precision
<Policy "mixed_float16">
>>> # Can optionally override layer to use float32 instead of mixed precision.
>>> layer2 = tf.keras.layers.Dense(10, dtype='float32')
>>> layer2.dtype_policy
<Policy "float32">
>>> # Set policy back to initial float32 for future examples.
>>> tf.keras.mixed_precision.set_global_policy('float32')
In the example above, passing `dtype='float32'` to the layer is equivalent to
passing `dtype=tf.keras.mixed_precision.Policy('float32')`. In general,
passing a dtype policy name to a layer is equivalent to passing the
corresponding policy, so it is never necessary to explicitly construct a
`Policy` object.
Note: `Model.compile` will automatically wrap an optimizer with a
`tf.keras.mixed_precision.LossScaleOptimizer` if you use the `'mixed_float16'`
policy. If you use a custom training loop instead of calling `Model.compile`,
you should explicitly use a `tf.keras.mixed_precision.LossScaleOptimizer` to
avoid numeric underflow with float16.
### How a layer uses its policy's compute dtype
A layer casts its inputs to its compute dtype. This causes the layer's
computations and output to also be in the compute dtype. For example:
>>> x = tf.ones((4, 4, 4, 4), dtype='float64')
>>> # `layer`'s policy defaults to float32.
>>> layer = tf.keras.layers.Conv2D(filters=4, kernel_size=2)
>>> layer.compute_dtype # Equivalent to layer.dtype_policy.compute_dtype
'float32'
>>> # `layer` casts its inputs to its compute dtype and does computations in
>>> # that dtype.
>>> y = layer(x)
>>> y.dtype
tf.float32
Note that the base `tf.keras.layers.Layer` class inserts the casts. If
subclassing your own layer, you do not have to insert any casts.
Currently, only tensors in the first argument to the layer's `call` method are
casted (although this will likely be changed in a future minor release). For
example:
>>> class MyLayer(tf.keras.layers.Layer):
... # Bug! `b` will not be casted.
... def call(self, a, b):
... return a + 1., b + 1.
>>> a = tf.constant(1., dtype="float32")
>>> b = tf.constant(1., dtype="float32")
>>> layer = MyLayer(dtype="float64")
>>> x, y = layer(a, b)
>>> x.dtype
tf.float64
>>> y.dtype
tf.float32
If writing your own layer with multiple inputs, you should either explicitly
cast other tensors to `self.compute_dtype` in `call` or accept all tensors in
the first argument as a list.
The casting only occurs in TensorFlow 2. If
`tf.compat.v1.disable_v2_behavior()` has been called, you can enable the
casting behavior with `tf.compat.v1.keras.layers.enable_v2_dtype_behavior()`.
### How a layer uses its policy's variable dtype
The default dtype of variables created by `tf.keras.layers.Layer.add_weight`
is the layer's policy's variable dtype.
If a layer's compute and variable dtypes differ, `add_weight` will wrap
floating-point variables with a special wrapper called an `AutoCastVariable`.
`AutoCastVariable` is identical to the original variable except it casts
itself to the layer's compute dtype when used within `Layer.call`. This means
if you are writing a layer, you do not have to explicitly cast the variables
to the layer's compute dtype. For example:
>>> class SimpleDense(tf.keras.layers.Layer):
...
... def build(self, input_shape):
... # With mixed precision, self.kernel is a float32 AutoCastVariable
... self.kernel = self.add_weight('kernel', (input_shape[-1], 10))
...
... def call(self, inputs):
... # With mixed precision, self.kernel will be casted to float16
... return tf.linalg.matmul(inputs, self.kernel)
...
>>> layer = SimpleDense(dtype='mixed_float16')
>>> y = layer(tf.ones((10, 10)))
>>> y.dtype
tf.float16
>>> layer.kernel.dtype
tf.float32
A layer author can prevent a variable from being wrapped with an
`AutoCastVariable` by passing `experimental_autocast=False` to `add_weight`,
which is useful if the float32 value of the variable must be accessed within
the layer.
### How to write a layer that supports mixed precision and float64.
For the most part, layers will automatically support mixed precision and
float64 without any additional work, due to the fact the base layer
automatically casts inputs, creates variables of the correct type, and in the
case of mixed precision, wraps variables with `AutoCastVariables`.
The primary case where you need extra work to support mixed precision or
float64 is when you create a new tensor, such as with `tf.ones` or
`tf.random.normal`, In such cases, you must create the tensor of the correct
dtype. For example, if you call `tf.random.normal`, you must pass the compute
dtype, which is the dtype the inputs have been casted to:
>>> class AddRandom(tf.keras.layers.Layer):
...
... def call(self, inputs):
... # We must pass `dtype=inputs.dtype`, otherwise a TypeError may
... # occur when adding `inputs` to `rand`.
... rand = tf.random.normal(shape=inputs.shape, dtype=inputs.dtype)
... return inputs + rand
>>> layer = AddRandom(dtype='mixed_float16')
>>> y = layer(x)
>>> y.dtype
tf.float16
If you did not pass `dtype=inputs.dtype` to `tf.random.normal`, a
`TypeError` would have occurred. This is because the `tf.random.normal`'s
dtype defaults to `"float32"`, but the input dtype is float16. You cannot add
a float32 tensor with a float16 tensor.
"""
def __init__(self, name):
if isinstance(name, dtypes.DType):
raise TypeError("'name' must be a string, not a DType. "
"Instead, pass DType.name. Got: %s" % (name.name,))
elif not isinstance(name, str):
raise TypeError("'name' must be a string, but got: %s" % (name,))
self._name = name
self._compute_dtype, self._variable_dtype = self._parse_name(name)
if name in ('mixed_float16', 'mixed_bloat16'):
device_compatibility_check.log_device_compatibility_check(name)
def _parse_name(self, name):
"""Parses a Policy name into a compute and variable dtype.
Args:
name: The name of the policy:
Returns:
The (compute_dtype, variable_dtype) pair.
"""
if name.endswith('_float32_vars'):
error_msg = ('Policies ending in \'_float32_vars\' have been removed '
'from TensorFlow.')
if name in ('infer_float32_vars', 'infer_with_float32_vars'):
error_msg += (' Please use the \'mixed_float16\' or \'mixed_bfloat16\' '
'policy instead.')
elif name == 'float16_with_float32_vars':
error_msg += (' Please use the \'mixed_float16\' policy instead.')
elif name == 'bfloat16_with_float32_vars':
error_msg += (' Please use the \'mixed_bfloat16\' policy instead.')
error_msg += ' Got policy name: \'%s\'' % name
raise ValueError(error_msg)
if name == 'mixed_float16':
return 'float16', 'float32'
elif name == 'mixed_bfloat16':
return 'bfloat16', 'float32'
elif name == '_infer':
# The "_infer" policy exists only for compatibility with TF 1, where
# "_infer" is the default. The behavior matches the behavior of TF 1's
# behavior before policies were introduced. With "_infer", the computation
# and variable dtype are inferred from the first input the first time the
# layer is called. Once the layer is called for the first time, the
# layer's policy will change to the dtype of the first input, and it will
# no longer have the "_infer" policy.
#
# The infer policy should be considered an implementation detail and may
# be removed in the future.
return None, None
try:
dtype = dtypes.as_dtype(name).name
except TypeError:
error = ("Cannot convert value %s to a mixed precision Policy. "
"Valid policies include 'mixed_float16', 'mixed_bfloat16', "
"and the name of any dtype such as 'float32'." % (name,))
raise ValueError(error)
return dtype, dtype
@property
def variable_dtype(self):
"""The variable dtype of this policy.
This is the dtype layers will create their variables in, unless a layer
explicitly chooses a different dtype. If this is different than
`Policy.compute_dtype`, Layers will cast variables to the compute dtype to
avoid type errors.
Variable regularizers are run in the variable dtype, not the compute dtype.
Returns:
The variable dtype of this policy, as a string.
"""
return self._variable_dtype
@property
def compute_dtype(self):
"""The compute dtype of this policy.
This is the dtype layers will do their computations in. Typically layers
output tensors with the compute dtype as well.
Note that even if the compute dtype is float16 or bfloat16, hardware devices
may not do individual adds, multiplies, and other fundamental operations in
float16 or bfloat16, but instead may do some of them in float32 for numeric
stability. The compute dtype is the dtype of the inputs and outputs of the
TensorFlow ops that the layer executes. Internally, many TensorFlow ops will
do certain internal calculations in float32 or some other device-internal
intermediate format with higher precision than float16/bfloat16, to increase
numeric stability.
For example, a `tf.keras.layers.Dense` layer, when run on a GPU with a
float16 compute dtype, will pass float16 inputs to `tf.linalg.matmul`. But,
`tf.linalg.matmul` will do use float32 intermediate math. The performance
benefit of float16 is still apparent, due to increased memory bandwidth and
the fact modern GPUs have specialized hardware for computing matmuls on
float16 inputs while still keeping intermediate computations in float32.
Returns:
The compute dtype of this policy, as a string.
"""
return self._compute_dtype
@property
def name(self):
"""Returns the name of this policy."""
return self._name
def __repr__(self):
return '<Policy "%s">' % self._name
def get_config(self):
return {'name': self.name}
@classmethod
def from_config(cls, config, custom_objects=None):
del custom_objects
if 'loss_scale' in config:
config = config.copy()
# Policy.get_config in TensorFlow 2.3 and below had a loss_scale. We
# silently drop it.
del config['loss_scale']
return cls(**config)
class PolicyV1(Policy):
"""A deprecated dtype policy for a Keras layer.
Warning: This class is now deprecated and will be removed soon. Please use the
non-experimental class `tf.keras.mixed_precision.Policy` instead.
The difference between this class and the non-experimental class is that this
class has a `loss_scale` field and the non-experimental class does not. The
loss scale is only used by `tf.keras.Model.compile`, which automatically wraps
the optimizer with a `LossScaleOptimizer` if the optimizer is not already a
`LossScaleOptimizer`. For the non-experimental Policy class, `Model.compile`
instead wraps the optimizer with a `LossScaleOptimizer` if `Policy.name` is
"mixed_float16".
When deserializing objects with an experimental policy using functions like
`tf.keras.utils.deserialize_keras_object`, the policy will be deserialized as
the non-experimental `tf.keras.mixed_precision.Policy`, and the loss scale
will silently be dropped. This is so that SavedModels that are generated
with an experimental policy can be restored after the experimental policy is
removed.
"""
def __init__(self, name, loss_scale='auto'):
"""Constructs the policy.
The `name` argument determines the compute and variable dtype, the default
loss scale, and has no additional effect on the Policy. The compute and
variable dtypes can only be specified through `name`, and cannot be
specified directly.
Args:
name: A string. Can be one of the following values:
* Any dtype name, such as 'float32' or 'float64'. Both the variable and
compute dtypes will be that dtype.
* 'mixed_float16' or 'mixed_bfloat16': The compute dtype is float16 or
bfloat16, while the variable dtype is float32. With 'mixed_float16',
a dynamic loss scale is used. These policies are used for mixed
precision training.
loss_scale: A `tf.compat.v1.mixed_precision.LossScale`, an int (which
uses a `FixedLossScale`), the string "dynamic" (which uses a
`DynamicLossScale`), or None (which uses no loss scale). Defaults to
`"auto"`. In the `"auto"` case: 1) if `name` is `"mixed_float16"`, then
use `loss_scale="dynamic"`. 2) otherwise, do not use a loss scale. Only
`tf.keras.Model`s, not layers, use the loss scale, and it is only used
during `Model.fit`, `Model.train_on_batch`, and other similar methods.
"""
super(PolicyV1, self).__init__(name)
if loss_scale == 'auto':
loss_scale = 'dynamic' if name == 'mixed_float16' else None
self._using_default_loss_scale = True
else:
self._using_default_loss_scale = False
if loss_scale and self._compute_dtype not in (None, 'float16'):
tf_logging.warning(
'Creating a Policy with a loss scale is only useful for '
'float16 policies. You passed loss_scale=%r for policy '
'%s. Consider not passing any loss_scale instead.' %
(loss_scale, name))
self._loss_scale = keras_loss_scale_module.get(loss_scale)
@property
def loss_scale(self):
"""Returns the loss scale of this Policy.
Returns:
A `tf.compat.v1.mixed_precision.experimental.LossScale`, or None.
"""
return self._loss_scale
def __repr__(self):
return '<PolicyV1 "%s", loss_scale=%s>' % (self._name, self.loss_scale)
def get_config(self):
config = {
'name': self.name
}
if not self._using_default_loss_scale:
# We only include the loss scale if the default loss scale is not used.
# This allows us to change the loss scale config format without breaking
# users who use the default loss scale.
config['loss_scale'] = keras_loss_scale_module.serialize(self.loss_scale)
return config
@classmethod
def from_config(cls, config, custom_objects=None):
if 'loss_scale' in config and isinstance(config['loss_scale'], dict):
config = config.copy()
config['loss_scale'] = keras_loss_scale_module.deserialize(
config['loss_scale'], custom_objects=custom_objects)
return cls(**config)
# The current global policy in effect. If None, it means the current value of
# floatx should be used as the policy if the V2 dtype behavior is enabled,
# or "_infer" otherwise.
# TODO(reedwm): Make this thread local?
_global_policy = None
def global_policy():
"""Returns the global dtype policy.
The global policy is the default `tf.keras.mixed_precision.Policy` used for
layers, if no policy is passed to the layer constructor. If no policy has been
set with `keras.mixed_precision.set_global_policy`, this will return a policy
constructed from `tf.keras.backend.floatx()` (floatx defaults to float32).
>>> tf.keras.mixed_precision.global_policy()
<Policy "float32">
>>> tf.keras.layers.Dense(10).dtype_policy # Defaults to the global policy
<Policy "float32">
If TensorFlow 2 behavior has been disabled with
`tf.compat.v1.disable_v2_behavior()`, this will instead return a special
"_infer" policy which infers the dtype from the dtype of the first input the
first time the layer is called. This behavior matches the behavior that
existed in TensorFlow 1.
See `tf.keras.mixed_precision.Policy` for more information on policies.
Returns:
The global Policy.
"""
if _global_policy is None:
if base_layer_utils.v2_dtype_behavior_enabled():
return Policy(backend.floatx())
else:
return Policy('_infer')
return _global_policy
def _check_if_mixed_precision_graph_rewrite_is_enabled(policy):
if mixed_precision_global_state.is_mixed_precision_graph_rewrite_enabled():
raise ValueError(
'The global dtype policy cannot be set to "{policy.name}", because the '
'mixed precision graph rewrite has already been enabled.\n'
'At most, one of the following can be called:\n\n'
' 1. tf.compat.v1.train.enable_mixed_precision_graph_rewrite() '
'(You called this first)\n'
' 2. tf.keras.mixed_precision.experimental.set_global_policy() with a '
'mixed precision policy (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 supports Eager execution and is more '
'customizable.'.format(policy=policy))
def set_global_policy(policy):
"""Sets the global dtype policy.
The global policy is the default `tf.keras.mixed_precision.Policy` used for
layers, if no policy is passed to the layer constructor.
>>> tf.keras.mixed_precision.set_global_policy('mixed_float16')
>>> tf.keras.mixed_precision.global_policy()
<Policy "mixed_float16">
>>> tf.keras.layers.Dense(10).dtype_policy
<Policy "mixed_float16">
>>> # Global policy is not used if a policy is directly passed to constructor
>>> tf.keras.layers.Dense(10, dtype='float64').dtype_policy
<Policy "float64">
>>> tf.keras.mixed_precision.set_global_policy('float32')
If no global policy is set, layers will instead default to a Policy
constructed from `tf.keras.backend.floatx()`.
To use mixed precision, the global policy should be set to `'mixed_float16'`
or `'mixed_bfloat16'`, so that every layer uses a 16-bit compute dtype and
float32 variable dtype by default.
Only floating point policies can be set as the global policy, such as
`'float32'` and `'mixed_float16'`. Non-floating point policies such as
`'int32'` and `'complex64'` cannot be set as the global policy because most
layers do not support such policies.
See `tf.keras.mixed_precision.Policy` for more information.
Args:
policy: A Policy, or a string that will be converted to a Policy. Can also
be None, in which case the global policy will be constructed from
`tf.keras.backend.floatx()`
"""
global _global_policy
if not base_layer_utils.v2_dtype_behavior_enabled():
raise ValueError('The global policy can only be set in TensorFlow 2 or if '
'V2 dtype behavior has been set. To enable V2 dtype '
'behavior, call '
'"tf.compat.v1.keras.layers.enable_v2_dtype_behavior()"')
if policy is not None and not isinstance(policy, Policy):
policy = Policy(policy)
is_mixed_policy = (policy is not None and
policy.compute_dtype != policy.variable_dtype)
if is_mixed_policy:
_check_if_mixed_precision_graph_rewrite_is_enabled(policy)
if (policy is not None and policy.compute_dtype is not None and
not dtypes.as_dtype(policy.compute_dtype).is_floating):
raise ValueError('set_global_policy can only be used to set the global '
'policy to floating-point policies, such as "float32" and '
'"mixed_float16", but got policy: %s'
% (policy.name,))
_global_policy = policy
mixed_precision_global_state.set_using_mixed_precision_policy(is_mixed_policy)
# TODO(reedwm): Make this thread local
@contextlib.contextmanager
def policy_scope(policy):
"""A context manager that sets the global Policy under it.
Args:
policy: A Policy, or a string that will be converted to a Policy..
Yields:
Nothing.
"""
old_policy = _global_policy
try:
set_global_policy(policy)
yield
finally:
set_global_policy(old_policy)
def _is_convertible_to_dtype(dtype):
try:
dtypes.as_dtype(dtype)
return True
except TypeError:
return False
def _policy_equivalent_to_dtype(policy):
"""Returns True if the Policy is equivalent to a single dtype.
A policy is equivalent to a single dtype if the policy's compute and variable
dtypes are the same and the policy's type is Policy and not a subclass of
Policy (such as PolicyV1).
The "_infer" policy is considered equivalent to a single dtype.
Args:
policy: A Policy.
Returns:
True, if the policy is equivalent to a single dtype.
"""
# We use type() instead of isinstance because a subclass of Policy is never
# equivalent to a dtype.
return (type(policy) == Policy and # pylint: disable=unidiomatic-typecheck
list(policy.get_config().keys()) == ['name'] and
(policy.name == '_infer' or _is_convertible_to_dtype(policy.name)))
def serialize(policy):
if _policy_equivalent_to_dtype(policy):
# We return either None or the policy name for compatibility with older
# versions of Keras. If the policy name is returned, it is a dtype string
# such as 'float32'.
return None if policy.name == '_infer' else policy.name
return generic_utils.serialize_keras_object(policy)
def deserialize(config, custom_objects=None):
if isinstance(config, str) and _is_convertible_to_dtype(config):
return Policy(config)
if config is None:
return Policy('_infer')
module_objects = {'Policy': Policy, 'PolicyV1': Policy}
return generic_utils.deserialize_keras_object(
config,
module_objects=module_objects,
custom_objects=custom_objects,
printable_module_name='dtype policy')
tensorflow/python/keras/mixed_precision/test_util.py
+225
View File
@@ -0,0 +1,225 @@
# Copyright 2019 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.
# ==============================================================================
"""Contains testing utilities related to mixed precision."""
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_conversion
from tensorflow.python.keras import regularizers
from tensorflow.python.keras.engine import base_layer
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import cond
from tensorflow.python.ops import custom_gradient
from tensorflow.python.ops import math_ops
from tensorflow.python.util import nest
def create_identity_with_grad_check_fn(expected_gradient, expected_dtype=None):
"""Returns a function that asserts it's gradient has a certain value.
This serves as a hook to assert intermediate gradients have a certain value.
This returns an identity function. The identity's gradient function is also
the identity function, except it asserts that the gradient equals
`expected_gradient` and has dtype `expected_dtype`.
Args:
expected_gradient: The gradient function asserts that the gradient is this
value.
expected_dtype: The gradient function asserts the gradient has this dtype.
Returns:
An identity function whose gradient function asserts the gradient has a
certain value.
"""
@custom_gradient.custom_gradient
def _identity_with_grad_check(x):
"""Function that asserts it's gradient has a certain value."""
x = array_ops.identity(x)
def grad(dx):
"""Gradient function that asserts the gradient has a certain value."""
if expected_dtype:
assert dx.dtype == expected_dtype, (
'dx.dtype should be %s but is: %s' % (expected_dtype, dx.dtype))
expected_tensor = tensor_conversion.convert_to_tensor_v2_with_dispatch(
expected_gradient, dtype=dx.dtype, name='expected_gradient'
)
# Control dependency is to ensure input is available. It's possible the
# dataset will throw a StopIteration to indicate there is no more data, in
# which case we don't want to run the assertion.
with ops.control_dependencies([x]):
assert_op = check_ops.assert_equal(dx, expected_tensor)
with ops.control_dependencies([assert_op]):
dx = array_ops.identity(dx)
return dx
return x, grad
# Keras sometimes has trouble serializing Lambda layers with a decorated
# function. So we define and return a non-decorated function.
def identity_with_grad_check(x):
return _identity_with_grad_check(x)
return identity_with_grad_check
def create_identity_with_nan_gradients_fn(have_nan_gradients):
"""Returns a function that optionally has NaN gradients.
This serves as a hook to introduce NaN gradients to a model. This returns an
identity function. The identity's gradient function will check if the boolean
tensor `have_nan_gradients` is True. If so, the gradient will be NaN.
Otherwise, the gradient will also be the identity.
Args:
have_nan_gradients: A scalar boolean tensor. If True, gradients will be NaN.
Otherwise, the gradient function is the identity function.
Returns:
An identity function whose gradient function will return NaNs, if
`have_nan_gradients` is True.
"""
@custom_gradient.custom_gradient
def _identity_with_nan_gradients(x):
"""Function whose gradient is NaN iff `have_nan_gradients` is True."""
x = array_ops.identity(x)
def grad(dx):
return cond.cond(
have_nan_gradients,
lambda: dx * float('NaN'),
lambda: dx
)
return x, grad
# Keras sometimes has trouble serializing Lambda layers with a decorated
# function. So we define and return a non-decorated function.
def identity_with_nan_gradients(x):
return _identity_with_nan_gradients(x)
return identity_with_nan_gradients
class AssertTypeLayer(base_layer.Layer):
"""A layer which asserts it's inputs are a certain type."""
def __init__(self, assert_type=None, **kwargs):
self._assert_type = (dtypes.as_dtype(assert_type).name if assert_type
else None)
super(AssertTypeLayer, self).__init__(**kwargs)
def assert_input_types(self, inputs):
"""Asserts `inputs` are of the correct type. Should be called in call()."""
if self._assert_type:
inputs_flattened = nest.flatten(inputs)
for inp in inputs_flattened:
assert inp.dtype.base_dtype == self._assert_type, (
'Input tensor has type %s which does not match assert type %s' %
(inp.dtype.name, self._assert_type))
class MultiplyLayer(AssertTypeLayer):
"""A layer which multiplies its input by a scalar variable."""
def __init__(self,
regularizer=None,
activity_regularizer=None,
use_operator=False,
var_name='v',
**kwargs):
"""Initializes the MultiplyLayer.
Args:
regularizer: The weight regularizer on the scalar variable.
activity_regularizer: The activity regularizer.
use_operator: If True, add using the * operator. If False, add using
tf.multiply.
var_name: The name of the variable. It can be useful to pass a name other
than 'v', to test having the attribute name (self.v) being different
from the variable name.
**kwargs: Passed to AssertTypeLayer constructor.
"""
self._regularizer = regularizer
if isinstance(regularizer, dict):
self._regularizer = regularizers.deserialize(regularizer,
custom_objects=globals())
self._activity_regularizer = activity_regularizer
if isinstance(activity_regularizer, dict):
self._activity_regularizer = regularizers.deserialize(
activity_regularizer, custom_objects=globals())
self._use_operator = use_operator
self._var_name = var_name
super(MultiplyLayer, self).__init__(
activity_regularizer=self._activity_regularizer, **kwargs)
def build(self, _):
self.v = self.add_weight(
self._var_name, (), initializer='ones', regularizer=self._regularizer)
self.built = True
def call(self, inputs):
self.assert_input_types(inputs)
return self._multiply(inputs, self.v)
def _multiply(self, x, y):
if self._use_operator:
return x * y
else:
return math_ops.multiply(x, y)
def get_config(self):
config = super(MultiplyLayer, self).get_config()
config['regularizer'] = regularizers.serialize(self._regularizer)
config['activity_regularizer'] = regularizers.serialize(
self._activity_regularizer)
config['use_operator'] = self._use_operator
config['var_name'] = self._var_name
config['assert_type'] = self._assert_type
return config
class MultiplyLayerWithoutAutoCast(MultiplyLayer):
"""Same as MultiplyLayer, but does not use AutoCastVariables."""
def build(self, _):
dtype = self.dtype
if dtype in ('float16', 'bfloat16'):
dtype = 'float32'
self.v = self.add_weight(
'v', (),
initializer='ones',
dtype=dtype,
experimental_autocast=False,
regularizer=self._regularizer)
self.built = True
def call(self, inputs):
self.assert_input_types(inputs)
assert self.v.dtype in (dtypes.float32, dtypes.float64)
return self._multiply(inputs, math_ops.cast(self.v, inputs.dtype))
class IdentityRegularizer(regularizers.Regularizer):
def __call__(self, x):
assert x.dtype == dtypes.float32
return array_ops.identity(x)
def get_config(self):
return {}
class ReduceSumRegularizer(regularizers.Regularizer):
def __call__(self, x):
return math_ops.reduce_sum(x)
def get_config(self):
return {}