3030

pending completion

Build # 3030

Build Type

Pull #889

travis-ci

Committed by

web-flow

Commit Message

<a href="https://github.com/google/trax/commit/<a class=hub.com/google/trax/commit/<a class="double-link" href="https://git"><a class=hub.com/google/trax/commit/<a class="double-link" href="https://git"><a class=hub.com/google/trax/commit/<a class="double-link" href="https://git"><a class=hub.com/google/trax/commit/<a class="double-link" href="https://git"><a class=hub.com/google/trax/commit/7a33c669206bd5dfd8f61fd94f7600620fa4ee43">7a33c6692">&lt;a href=&quot;https://github.com/google/trax/commit/</a><a class="double-link" href="https://github.com/google/trax/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/google/trax/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/google/trax/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/google/trax/commit/7a33c669206bd5dfd8f61fd94f7600620fa4ee43">7a33c6692</a><a href="https://github.com/google/trax/commit/7a33c669206bd5dfd8f61fd94f7600620fa4ee43">&lt;a href=&quot;https://github.com/google/trax/commit/7a33c669206bd5dfd8f61fd94f7600620fa4ee43&quot;&gt;&amp;quot;&amp;gt;&amp;amp;lt;a href=&amp;amp;quot;https://github.com/google/trax/commit/&amp;lt;/a&amp;gt;&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://github.com/google/trax/commit/&amp;amp;lt;a class=&amp;amp;quot;double-link&amp;amp;quot; href=&amp;amp;quot;https://git&amp;quot;&amp;gt;&amp;amp;lt;a class=&amp;lt;/a&amp;gt;hub.com/google/trax/commit/7a33c669206bd5dfd8f61fd94f7600620fa4ee43&amp;quot;&amp;gt;7a33c6692&amp;lt;/a&amp;gt;&amp;lt;a href=&amp;quot;https://github.com/google/trax/commit/7a33c669206bd5dfd8f61fd94f7600620fa4ee43&amp;quot;&amp;gt;&amp;amp;lt;a href=&amp;amp;quot;https://github.com/google/trax/commit/7a33c669206bd5dfd8f61fd94f760062&amp;lt;/a&amp;gt</a>;0fa4ee43&quot;&gt;&amp;quot;&amp;gt;Merge &amp;lt;/a&amp;gt;&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://github.com/google/trax/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://github.com/google/trax/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://github.com/google/trax/commit/aafc0d824&quot;&gt;aafc0d824&lt;/a&gt;&lt;a href=&quot;https://github.com/google/trax/commit/7a33c669206bd5dfd8f61fd94f760... (continued)

Pull Request Pull Request #889: [TRAX] Backends are enums instead of strings.

Run Details

50 of 50 new or added lines in 15 files covered. (100.0%)

2252 of 9833 relevant lines covered (22.9%)

0.23 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

57.14

/trax/layers/metrics.py

# coding=utf-8
# Copyright 2020 The Trax Authors.
#
# 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.

# Lint as: python3
"""Trax layers for computing metrics (loss functions and evaluation metrics).

A metric layer takes three inputs:

  - model output: Batch of predicted values (typically vectors).
  - targets: Batch of target values (e.g., categories or vectors).
  - weights: Tensor that can assign different weights to different positions
    in the model output. One common use of weights is for masking -- assigning
    weight 0 to positions that correspond to padding in the input so that they
    don't affect metrics.

and returns a single scalar.

The `L2Loss` layer treats a batch as an unanalyzed tensor and computes an
elementwise-weighted loss.

Other metric layers take into account the items that make up a batch. For each
item in a batch, a raw metric value is computed by comparing (item-wise) the
model output to the target value. These item-wise values are then combined into
a single scalar for the batch by a weighted reduction function, typically
weighted mean. For example:

  - Accuracy: Treat model output as giving different strength/votes to the
    possible categories; measure the category prediction as correct (value 1)
    if `argmax(output) == target_category`, else as incorrect (value 0). The
    accuracy for the batch is then the weighted mean of these 1's and 0's.

  - Cross Entropy: Treat model output and target values as two probability
    distributions; measure the cross entropy of the model output relative to
    the (assumed true) target distribution. The scalar value for the batch is
    then the weighted mean of the item-wise cross-entropy values.

In deriving a single scalar for the batch, there is flexibility to use reducing
functions other than mean, for instance sum or a specialized sequence mean.
"""

import jax

from trax import fastmath
from trax import shapes
from trax.fastmath import numpy as jnp
from trax.layers import combinators as cb
from trax.layers import core
from trax.layers.base import Fn


def L2Loss():
  """Returns a layer that computes total L2 loss for one batch."""
  def f(model_output, targets, weights):  # pylint: disable=invalid-name
    """Returns elementwise-weighted L2 norm of `model_output - targets`.

    Args:
      model_output: Output from one batch, treated as an unanalyzed tensor.
      targets: Tensor of same shape as `model_output` containing element-wise
          target values.
      weights: Tensor of same shape as `model_output` and `targets`.
    """
    shapes.assert_same_shape(model_output, targets)
    shapes.assert_same_shape(targets, weights)
    l2 = weights * (model_output - targets)**2
    return jnp.sum(l2) / jnp.sum(weights)
  return Fn('L2Loss', f)


def Accuracy():
  """Returns a layer that computes mean category prediction accuracy."""
  return cb.Serial(_Accuracy(), _WeightedMean(),
                   name='Accuracy', sublayers_to_print=[])


def SequenceAccuracy():
  """Returns a layer that computes mean sequence prediction accuracy."""
  return cb.Serial(_Accuracy(), _WeightedSequenceMean(),
                   name='SequenceAccuracy', sublayers_to_print=[])


def CrossEntropyLoss():
  """Returns a layer that computes mean prediction-target cross entropy."""
  return cb.Serial(_CrossEntropy(), _WeightedMean(),
                   name='CrossEntropyLoss', sublayers_to_print=[])


def CrossEntropySum():
  """Returns a layer that computes sum of prediction-target cross entropies."""
  return cb.Serial(_CrossEntropy(), WeightedSum(),
                   name='CrossEntropySum', sublayers_to_print=[])


def SumOfWeights():
  """Returns a layer that computes sum of weights."""
  return cb.Serial(
      cb.Drop(),  # Drop inputs.
      cb.Drop(),  # Drop targets.
      core.Sum(axis=None)  # Sum weights.
  )
# pylint: enable=no-value-for-parameter


def _Accuracy(axis=-1):
  """Returns a layer that scores predicted versus target category."""
  def f(model_output, target_category):  # pylint: disable=invalid-name
    predicted_category = jnp.argmax(model_output, axis=axis)
    # TODO(pkozakowski): This assertion breaks some tests. Fix and uncomment.
    # shapes.assert_same_shape(predicted_category, target_category)
    return jnp.equal(predicted_category, target_category).astype(jnp.float32)
  return Fn('_Accuracy', f)


def _CrossEntropy():
  """Returns a layer that computes prediction-target cross entropies."""
  def f(model_output, target_category):  # pylint: disable=invalid-name
    # TODO(pkozakowski): This assertion breaks some tests. Fix and uncomment.
    # shapes.assert_shape_equals(target_category, model_output.shape[:-1])
    target_distribution = one_hot(target_category, model_output.shape[-1])
    return -1.0 * jnp.sum(model_output * target_distribution, axis=-1)
  return Fn('_CrossEntropy', f)


def _WeightedMean():
  """Returns a layer that computes a weighted mean of the given values."""
  def f(values, weights):  # pylint: disable=invalid-name
    return jnp.sum(values * weights) / jnp.sum(weights)
  return Fn('_WeightedMean', f)


def WeightedSum():
  """Returns a layer that computes a weighted sum of the given values."""
  def f(values, weights):  # pylint: disable=invalid-name
    return jnp.sum(values * weights)
  return Fn('WeightedSum', f)


def _WeightedSequenceMean():
  """Returns a layer that computes a weighted sequence accuracy mean."""
  def f(values, weights):  # pylint: disable=invalid-name
    # This function assumes weights are 0 or 1.
    # Then compute 1: not-correct, 0: correct or masked
    not_correct = (1.0 - values) * weights
    axis_to_sum = list(range(1, len(not_correct.shape)))
    # Summing not-correct on all axes but batch. We're summing 0s and 1s,
    # so the sum is 0 if it's all 0 and >=1 in all other cases.
    not_correct_seq = jnp.sum(not_correct, axis=axis_to_sum)
    # Sequence is correct if not_correct_seq is 0, reverting here.
    correct_seq = 1.0 - jnp.minimum(1.0, not_correct_seq)
    return jnp.mean(correct_seq)  # Mean over batch.
  return Fn('_WeightedSequenceMean', f)


# TODO(jonni): Figure out the right name and home for this function.
def one_hot(x, n_categories, dtype=jnp.float32):  # pylint: disable=invalid-name
  """Makes a one-hot array (n+1 dims) from an int-categorical array (n dims)."""
  indices_less_than_n = jnp.arange(n_categories)
  if fastmath.is_backend(fastmath.Backend.JAX):
    # Work around a jax broadcasting issue.
    indices_less_than_n = jax.lax.tie_in(x, indices_less_than_n)
  return jnp.array(x[..., jnp.newaxis] == indices_less_than_n, dtype)

1	# coding=utf-8
2	# Copyright 2020 The Trax Authors.
3	#
4	# Licensed under the Apache License, Version 2.0 (the "License");
5	# you may not use this file except in compliance with the License.
6	# You may obtain a copy of the License at
7	#
8	# http://www.apache.org/licenses/LICENSE-2.0
9	#
10	# Unless required by applicable law or agreed to in writing, software
11	# distributed under the License is distributed on an "AS IS" BASIS,
12	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	# See the License for the specific language governing permissions and
14	# limitations under the License.
15
16	# Lint as: python3
17	"""Trax layers for computing metrics (loss functions and evaluation metrics).	1✔
18
19	A metric layer takes three inputs:
20
21	- model output: Batch of predicted values (typically vectors).
22	- targets: Batch of target values (e.g., categories or vectors).
23	- weights: Tensor that can assign different weights to different positions
24	in the model output. One common use of weights is for masking -- assigning
25	weight 0 to positions that correspond to padding in the input so that they
26	don't affect metrics.
27
28	and returns a single scalar.
29
30	The `L2Loss` layer treats a batch as an unanalyzed tensor and computes an
31	elementwise-weighted loss.
32
33	Other metric layers take into account the items that make up a batch. For each
34	item in a batch, a raw metric value is computed by comparing (item-wise) the
35	model output to the target value. These item-wise values are then combined into
36	a single scalar for the batch by a weighted reduction function, typically
37	weighted mean. For example:
38
39	- Accuracy: Treat model output as giving different strength/votes to the
40	possible categories; measure the category prediction as correct (value 1)
41	if `argmax(output) == target_category`, else as incorrect (value 0). The
42	accuracy for the batch is then the weighted mean of these 1's and 0's.
43
44	- Cross Entropy: Treat model output and target values as two probability
45	distributions; measure the cross entropy of the model output relative to
46	the (assumed true) target distribution. The scalar value for the batch is
47	then the weighted mean of the item-wise cross-entropy values.
48
49	In deriving a single scalar for the batch, there is flexibility to use reducing
50	functions other than mean, for instance sum or a specialized sequence mean.
51	"""
52
53	import jax	1✔
54
55	from trax import fastmath	1✔
56	from trax import shapes	1✔
57	from trax.fastmath import numpy as jnp	1✔
58	from trax.layers import combinators as cb	1✔
59	from trax.layers import core	1✔
60	from trax.layers.base import Fn	1✔
61
62
63	def L2Loss():	1✔
64	"""Returns a layer that computes total L2 loss for one batch."""
65	def f(model_output, targets, weights): # pylint: disable=invalid-name	×
66	"""Returns elementwise-weighted L2 norm of `model_output - targets`.
67
68	Args:
69	model_output: Output from one batch, treated as an unanalyzed tensor.
70	targets: Tensor of same shape as `model_output` containing element-wise
71	target values.
72	weights: Tensor of same shape as `model_output` and `targets`.
73	"""
74	shapes.assert_same_shape(model_output, targets)	×
75	shapes.assert_same_shape(targets, weights)	×
76	l2 = weights * (model_output - targets)**2	×
77	return jnp.sum(l2) / jnp.sum(weights)	×
78	return Fn('L2Loss', f)	×
79
80
81	def Accuracy():	1✔
82	"""Returns a layer that computes mean category prediction accuracy."""
83	return cb.Serial(_Accuracy(), _WeightedMean(),	1✔
84	name='Accuracy', sublayers_to_print=[])
85
86
87	def SequenceAccuracy():	1✔
88	"""Returns a layer that computes mean sequence prediction accuracy."""
89	return cb.Serial(_Accuracy(), _WeightedSequenceMean(),	1✔
90	name='SequenceAccuracy', sublayers_to_print=[])
91
92
93	def CrossEntropyLoss():	1✔
94	"""Returns a layer that computes mean prediction-target cross entropy."""
95	return cb.Serial(_CrossEntropy(), _WeightedMean(),	1✔
96	name='CrossEntropyLoss', sublayers_to_print=[])
97
98
99	def CrossEntropySum():	1✔
100	"""Returns a layer that computes sum of prediction-target cross entropies."""
101	return cb.Serial(_CrossEntropy(), WeightedSum(),	×
102	name='CrossEntropySum', sublayers_to_print=[])
103
104
105	def SumOfWeights():	1✔
106	"""Returns a layer that computes sum of weights."""
107	return cb.Serial(	1✔
108	cb.Drop(), # Drop inputs.
109	cb.Drop(), # Drop targets.
110	core.Sum(axis=None) # Sum weights.
111	)
112	# pylint: enable=no-value-for-parameter
113
114
115	def _Accuracy(axis=-1):	1✔
116	"""Returns a layer that scores predicted versus target category."""
117	def f(model_output, target_category): # pylint: disable=invalid-name	1✔
118	predicted_category = jnp.argmax(model_output, axis=axis)	×
119	# TODO(pkozakowski): This assertion breaks some tests. Fix and uncomment.
120	# shapes.assert_same_shape(predicted_category, target_category)
121	return jnp.equal(predicted_category, target_category).astype(jnp.float32)	×
122	return Fn('_Accuracy', f)	1✔
123
124
125	def _CrossEntropy():	1✔
126	"""Returns a layer that computes prediction-target cross entropies."""
127	def f(model_output, target_category): # pylint: disable=invalid-name	1✔
128	# TODO(pkozakowski): This assertion breaks some tests. Fix and uncomment.
129	# shapes.assert_shape_equals(target_category, model_output.shape[:-1])
130	target_distribution = one_hot(target_category, model_output.shape[-1])	×
131	return -1.0 * jnp.sum(model_output * target_distribution, axis=-1)	×
132	return Fn('_CrossEntropy', f)	1✔
133
134
135	def _WeightedMean():	1✔
136	"""Returns a layer that computes a weighted mean of the given values."""
137	def f(values, weights): # pylint: disable=invalid-name	1✔
138	return jnp.sum(values * weights) / jnp.sum(weights)	×
139	return Fn('_WeightedMean', f)	1✔
140
141
142	def WeightedSum():	1✔
143	"""Returns a layer that computes a weighted sum of the given values."""
144	def f(values, weights): # pylint: disable=invalid-name	×
145	return jnp.sum(values * weights)	×
146	return Fn('WeightedSum', f)	×
147
148
149	def _WeightedSequenceMean():	1✔
150	"""Returns a layer that computes a weighted sequence accuracy mean."""
151	def f(values, weights): # pylint: disable=invalid-name	1✔
152	# This function assumes weights are 0 or 1.
153	# Then compute 1: not-correct, 0: correct or masked
154	not_correct = (1.0 - values) * weights	×
155	axis_to_sum = list(range(1, len(not_correct.shape)))	×
156	# Summing not-correct on all axes but batch. We're summing 0s and 1s,
157	# so the sum is 0 if it's all 0 and >=1 in all other cases.
158	not_correct_seq = jnp.sum(not_correct, axis=axis_to_sum)	×
159	# Sequence is correct if not_correct_seq is 0, reverting here.
160	correct_seq = 1.0 - jnp.minimum(1.0, not_correct_seq)	×
161	return jnp.mean(correct_seq) # Mean over batch.	×
162	return Fn('_WeightedSequenceMean', f)	1✔
163
164
165	# TODO(jonni): Figure out the right name and home for this function.
166	def one_hot(x, n_categories, dtype=jnp.float32): # pylint: disable=invalid-name	1✔
167	"""Makes a one-hot array (n+1 dims) from an int-categorical array (n dims)."""
168	indices_less_than_n = jnp.arange(n_categories)	×
169	if fastmath.is_backend(fastmath.Backend.JAX):	×
170	# Work around a jax broadcasting issue.
171	indices_less_than_n = jax.lax.tie_in(x, indices_less_than_n)	×
172	return jnp.array(x[..., jnp.newaxis] == indices_less_than_n, dtype)	×

google / trax / 3030

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous