google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research 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
"""Tests for the scalarization optimizers."""
17

18
from absl.testing import parameterized
19
import tensorflow.compat.v1 as tf
20

21
from yoto.optimizers import scalarization
22
import yoto.problems as problems
23

24

25
class ProblemWithConstantLosses(problems.Problem):
26

27
  def __init__(self, losses_values):
28
    self._losses_values = losses_values
29
    self._dummy = tf.Variable(0.)
30

31
  def losses_and_metrics(self, inputs, inputs_extra=None, training=False):
32
    """Map the inputs to a {loss_name: loss_tensor} dictionary."""
33
    del inputs
34
    del inputs_extra
35
    losses = {key: value + 0 * self._dummy
36
              for key, value in self._losses_values.items()}
37
    return losses, {}
38

39
  @property
40
  def losses_keys(self):
41
    return tuple(sorted(self._losses_values.keys()))
42

43
  def initialize_model(self):
44
    pass
45

46
  def module_spec(self):
47
    pass
48

49

50
class LinearlyScalarizedOptimizerTest(parameterized.TestCase,
51
                                      tf.test.TestCase):
52

53
  def test_check_weighted_value_on_constant_losses(self):
54
    weights = {"a": tf.constant(0.5),
55
               "b": tf.constant(0.3),
56
               "c": tf.constant(0.4)}
57
    losses = {"a": -15, "b": .4, "c": .3}
58
    optimizer = scalarization.LinearlyScalarizedOptimizer(
59
        problem=ProblemWithConstantLosses(losses), weights=weights)
60
    loss, _ = optimizer.compute_train_loss_and_update_op(
61
        inputs=dict(), base_optimizer=tf.train.GradientDescentOptimizer(0.))
62
    with self.cached_session() as session:
63
      session.run(tf.initializers.global_variables())
64
    self.assertAllClose(loss,
65
                        sum(weights[key] * losses[key] for key in weights))
66

67
  def test_exception_thrown_when_weights_is_of_invalid_type(self):
68
    losses = {"a": -15, "b": .4, "c": .3}
69
    # Should fail as `weights` is neither a dict nor in the enum.
70
    with self.assertRaises(TypeError):
71
      _ = scalarization.LinearlyScalarizedOptimizer(
72
          problem=ProblemWithConstantLosses(losses), weights=123)
73

74
  def test_throws_exception_when_weights_key_is_missing(self):
75
    losses = {"a": -15, "b": .4, "c": .3}
76
    weights = {"a": tf.constant(0.5),
77
               "b": tf.constant(0.3)}  # Misses the key "c".
78
    with self.assertRaises(ValueError):
79
      _ = scalarization.LinearlyScalarizedOptimizer(
80
          problem=ProblemWithConstantLosses(losses), weights=weights)
81

82
if __name__ == "__main__":
83
  tf.disable_eager_execution()
84
  tf.test.main()
85