Move the LossReduction class from tf to Keras.

PiperOrigin-RevId: 351654535
Change-Id: I405d4f79568c05ddeaa3755c3171960db633e6de
This commit is contained in:
Scott Zhu 2021-01-13 13:33:28 -08:00 committed by TensorFlower Gardener
parent 02e9e26b27
commit c4aab50762
11 changed files with 73 additions and 106 deletions

View File

@ -151,7 +151,6 @@ py_library(
"//tensorflow/python:variable_scope",
"//tensorflow/python/data",
"//tensorflow/python/ops/losses",
"//tensorflow/python/ops/losses:loss_reduction",
"//tensorflow/tools/docs:doc_controls",
],
)

View File

@ -220,7 +220,6 @@ from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import summary_ops_v2
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops.losses import loss_reduction
from tensorflow.python.ops.losses import losses_impl
from tensorflow.python.platform import tf_logging
from tensorflow.python.training.tracking import base as trackable
@ -286,7 +285,7 @@ def get_loss_reduction():
return reduce_util.ReduceOp.SUM
last_reduction = ops.get_default_graph()._last_loss_reduction # pylint: disable=protected-access
if (last_reduction == losses_impl.Reduction.SUM or
last_reduction == loss_reduction.ReductionV2.SUM):
last_reduction == "sum"): # Check for tf.keras.losses.Reduction.SUM
return reduce_util.ReduceOp.SUM
return reduce_util.ReduceOp.MEAN

View File

@ -60,6 +60,7 @@ from tensorflow.python.keras.distribute.strategy_combinations import tpu_strateg
from tensorflow.python.keras.engine import base_layer_utils
from tensorflow.python.keras.mixed_precision import policy
from tensorflow.python.keras.optimizer_v2 import gradient_descent as gradient_descent_keras
from tensorflow.python.keras.utils import losses_utils
from tensorflow.python.keras.utils import np_utils
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import check_ops
@ -67,7 +68,6 @@ from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import parsing_ops
from tensorflow.python.ops import variables
from tensorflow.python.ops.losses import loss_reduction
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.platform import test
from tensorflow.python.training import gradient_descent
@ -2068,9 +2068,9 @@ class TestDistributionStrategyWithKerasModels(test.TestCase,
],
mode=['graph', 'eager'],
reduction=[
loss_reduction.ReductionV2.AUTO,
loss_reduction.ReductionV2.SUM_OVER_BATCH_SIZE,
loss_reduction.ReductionV2.SUM
losses_utils.ReductionV2.AUTO,
losses_utils.ReductionV2.SUM_OVER_BATCH_SIZE,
losses_utils.ReductionV2.SUM
]))
def test_distribution_strategy_with_loss_reduction_types(
self, distribution, reduction):
@ -2573,7 +2573,7 @@ class TestDistributionStrategyWithMultipleAddLossAndMetricCalls(
optimizer=keras.optimizers.adam_v2.Adam(1e-4),
loss=keras.losses.SparseCategoricalCrossentropy(
from_logits=True,
reduction=loss_reduction.ReductionV2.SUM_OVER_BATCH_SIZE),
reduction=losses_utils.ReductionV2.SUM_OVER_BATCH_SIZE),
metrics=[
keras.metrics.SparseCategoricalAccuracy(),
keras.metrics.SparseCategoricalCrossentropy(from_logits=True),

View File

@ -38,10 +38,10 @@ from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor_spec
from tensorflow.python.keras.layers.preprocessing import string_lookup
from tensorflow.python.keras.optimizer_v2 import rmsprop
from tensorflow.python.keras.utils import losses_utils
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops.losses import loss_reduction
from tensorflow.python.platform import test
from tensorflow.python.training.server_lib import ClusterSpec
@ -167,7 +167,7 @@ class KPLTest(test.TestCase, parameterized.TestCase):
pred = model(batch_data, training=True)
loss = nn.compute_average_loss(
keras.losses.BinaryCrossentropy(
reduction=loss_reduction.ReductionV2.NONE)(labels, pred))
reduction=losses_utils.ReductionV2.NONE)(labels, pred))
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))

View File

@ -26,7 +26,7 @@ from tensorflow.python.keras import layers
from tensorflow.python.keras import losses
from tensorflow.python.keras import metrics
from tensorflow.python.keras import testing_utils
from tensorflow.python.ops.losses import loss_reduction
from tensorflow.python.keras.utils import losses_utils
from tensorflow.python.platform import test
from tensorflow.python.util import nest
@ -557,9 +557,9 @@ class TestMetricsCorrectnessSingleIO(keras_parameterized.TestCase):
@keras_parameterized.run_with_all_model_types(exclude_models=['sequential'])
@keras_parameterized.run_all_keras_modes(always_skip_v1=True)
@parameterized.parameters([
loss_reduction.ReductionV2.SUM_OVER_BATCH_SIZE,
loss_reduction.ReductionV2.AUTO,
loss_reduction.ReductionV2.SUM
losses_utils.ReductionV2.SUM_OVER_BATCH_SIZE,
losses_utils.ReductionV2.AUTO,
losses_utils.ReductionV2.SUM
])
class TestOutputLossMetrics(keras_parameterized.TestCase):
@ -615,25 +615,25 @@ class TestOutputLossMetrics(keras_parameterized.TestCase):
}
self.expected_fit_result = {
loss_reduction.ReductionV2.NONE:
losses_utils.ReductionV2.NONE:
sum_over_batch_size_fit_result,
loss_reduction.ReductionV2.SUM: {
losses_utils.ReductionV2.SUM: {
'loss': [198, 198],
'output_1_loss': [82, 82],
'output_2_loss': [116, 116],
},
loss_reduction.ReductionV2.AUTO:
losses_utils.ReductionV2.AUTO:
sum_over_batch_size_fit_result,
loss_reduction.ReductionV2.SUM_OVER_BATCH_SIZE:
losses_utils.ReductionV2.SUM_OVER_BATCH_SIZE:
sum_over_batch_size_fit_result,
}
# In the order: 'loss', 'output_1_loss', 'output_2_loss',
self.expected_batch_result = {
loss_reduction.ReductionV2.NONE: [144, 56, 88],
loss_reduction.ReductionV2.SUM: [198, 82, 116],
loss_reduction.ReductionV2.AUTO: [144, 56, 88],
loss_reduction.ReductionV2.SUM_OVER_BATCH_SIZE: [144, 56, 88],
losses_utils.ReductionV2.NONE: [144, 56, 88],
losses_utils.ReductionV2.SUM: [198, 82, 116],
losses_utils.ReductionV2.AUTO: [144, 56, 88],
losses_utils.ReductionV2.SUM_OVER_BATCH_SIZE: [144, 56, 88],
}
# 2 + 12 + 36 + 80 + 150 = 280
@ -675,7 +675,7 @@ class TestOutputLossMetrics(keras_parameterized.TestCase):
})
expected_values = self.expected_batch_result[reduction]
if reduction == loss_reduction.ReductionV2.SUM:
if reduction == losses_utils.ReductionV2.SUM:
expected_values = self.expected_single_batch_result
self.assertAllClose(result, expected_values)
@ -688,7 +688,7 @@ class TestOutputLossMetrics(keras_parameterized.TestCase):
'output_2': self.sample_weight_2,
})
expected_values = self.expected_batch_result[reduction]
if reduction == loss_reduction.ReductionV2.SUM:
if reduction == losses_utils.ReductionV2.SUM:
expected_values = self.expected_single_batch_result
self.assertAllClose(result, expected_values)

View File

@ -74,7 +74,6 @@ py_library(
":data_utils",
":io_utils",
"//tensorflow/python/keras:backend",
"//tensorflow/python/ops/losses:loss_reduction",
],
)

View File

@ -25,15 +25,59 @@ from tensorflow.python.keras.engine import keras_tensor
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.losses import loss_reduction
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.util.tf_export import keras_export
# TODO(joshl/psv): Update references to ReductionV2 to point to its
# new location.
ReductionV2 = loss_reduction.ReductionV2
keras_export('keras.losses.Reduction', v1=[])(loss_reduction.ReductionV2)
@keras_export('keras.losses.Reduction', v1=[])
class ReductionV2(object):
"""Types of loss reduction.
Contains the following values:
* `AUTO`: Indicates that the reduction option will be determined by the usage
context. For almost all cases this defaults to `SUM_OVER_BATCH_SIZE`. When
used with `tf.distribute.Strategy`, outside of built-in training loops such
as `tf.keras` `compile` and `fit`, we expect reduction value to be
`SUM` or `NONE`. Using `AUTO` in that case will raise an error.
* `NONE`: Weighted losses with one dimension reduced (axis=-1, or axis
specified by loss function). When this reduction type used with built-in
Keras training loops like `fit`/`evaluate`, the unreduced vector loss is
passed to the optimizer but the reported loss will be a scalar value.
* `SUM`: Scalar sum of weighted losses.
* `SUM_OVER_BATCH_SIZE`: Scalar `SUM` divided by number of elements in losses.
This reduction type is not supported when used with
`tf.distribute.Strategy` outside of built-in training loops like `tf.keras`
`compile`/`fit`.
You can implement 'SUM_OVER_BATCH_SIZE' using global batch size like:
```
with strategy.scope():
loss_obj = tf.keras.losses.CategoricalCrossentropy(
reduction=tf.keras.losses.Reduction.NONE)
....
loss = tf.reduce_sum(loss_obj(labels, predictions)) *
(1. / global_batch_size)
```
Please see the [custom training guide](
https://www.tensorflow.org/tutorials/distribute/custom_training) for more
details on this.
"""
AUTO = 'auto'
NONE = 'none'
SUM = 'sum'
SUM_OVER_BATCH_SIZE = 'sum_over_batch_size'
@classmethod
def all(cls):
return (cls.AUTO, cls.NONE, cls.SUM, cls.SUM_OVER_BATCH_SIZE)
@classmethod
def validate(cls, key):
if key not in cls.all():
raise ValueError('Invalid Reduction Key %s.' % key)
def remove_squeezable_dimensions(

View File

@ -31,12 +31,6 @@ py_library(
],
)
py_library(
name = "loss_reduction",
srcs = ["loss_reduction.py"],
srcs_version = "PY2AND3",
)
py_test(
name = "util_test",
size = "small",

View File

@ -1,68 +0,0 @@
# 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.
# ==============================================================================
"""The TF2 version of the enum keras.losses.Reduction."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
class ReductionV2(object):
"""Types of loss reduction.
Contains the following values:
* `AUTO`: Indicates that the reduction option will be determined by the usage
context. For almost all cases this defaults to `SUM_OVER_BATCH_SIZE`. When
used with `tf.distribute.Strategy`, outside of built-in training loops such
as `tf.keras` `compile` and `fit`, we expect reduction value to be
`SUM` or `NONE`. Using `AUTO` in that case will raise an error.
* `NONE`: Weighted losses with one dimension reduced (axis=-1, or axis
specified by loss function). When this reduction type used with built-in
Keras training loops like `fit`/`evaluate`, the unreduced vector loss is
passed to the optimizer but the reported loss will be a scalar value.
* `SUM`: Scalar sum of weighted losses.
* `SUM_OVER_BATCH_SIZE`: Scalar `SUM` divided by number of elements in losses.
This reduction type is not supported when used with
`tf.distribute.Strategy` outside of built-in training loops like `tf.keras`
`compile`/`fit`.
You can implement 'SUM_OVER_BATCH_SIZE' using global batch size like:
```
with strategy.scope():
loss_obj = tf.keras.losses.CategoricalCrossentropy(
reduction=tf.keras.losses.Reduction.NONE)
....
loss = tf.reduce_sum(loss_obj(labels, predictions)) *
(1. / global_batch_size)
```
Please see the [custom training guide](
https://www.tensorflow.org/tutorials/distribute/custom_training) for more
details on this.
"""
AUTO = 'auto'
NONE = 'none'
SUM = 'sum'
SUM_OVER_BATCH_SIZE = 'sum_over_batch_size'
@classmethod
def all(cls):
return (cls.AUTO, cls.NONE, cls.SUM, cls.SUM_OVER_BATCH_SIZE)
@classmethod
def validate(cls, key):
if key not in cls.all():
raise ValueError('Invalid Reduction Key %s.' % key)

View File

@ -1,6 +1,6 @@
path: "tensorflow.keras.losses.Reduction"
tf_class {
is_instance: "<class \'tensorflow.python.ops.losses.loss_reduction.ReductionV2\'>"
is_instance: "<class \'tensorflow.python.keras.utils.losses_utils.ReductionV2\'>"
is_instance: "<type \'object\'>"
member {
name: "AUTO"

View File

@ -1,6 +1,6 @@
path: "tensorflow.losses.Reduction"
tf_class {
is_instance: "<class \'tensorflow.python.ops.losses.loss_reduction.ReductionV2\'>"
is_instance: "<class \'tensorflow.python.keras.utils.losses_utils.ReductionV2\'>"
is_instance: "<type \'object\'>"
member {
name: "AUTO"