|
| 1 | +# Copyright 2021 Alibaba Group Holding Limited. All Rights Reserved. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | +# ============================================================================= |
| 15 | +"""Test for multiple optimizer.""" |
| 16 | + |
| 17 | +from __future__ import absolute_import |
| 18 | +from __future__ import division |
| 19 | +from __future__ import print_function |
| 20 | + |
| 21 | +import numpy as np |
| 22 | +import tensorflow as tf |
| 23 | +from tensorflow.python.platform import test |
| 24 | + |
| 25 | +import epl |
| 26 | +from test_utils import fix_randomness |
| 27 | + |
| 28 | + |
| 29 | +def get_nested_optimizer(): |
| 30 | + optimizer = tf.train.AdamOptimizer(learning_rate=0.01) |
| 31 | + optimizer = tf.contrib.estimator.clip_gradients_by_norm(optimizer, 0.5) |
| 32 | + optimizer = tf.contrib.opt.MovingAverageOptimizer( |
| 33 | + optimizer, |
| 34 | + average_decay=0.1, |
| 35 | + num_updates=2) |
| 36 | + return optimizer |
| 37 | + |
| 38 | +def _train_fn(bs, lr_decay=False, max_steps=5, optimizer=None, enable_epl=True): |
| 39 | + """Define train function.""" |
| 40 | + fix_randomness() |
| 41 | + num_x = np.random.randint(0, 10, (500, 20)).astype(dtype=np.float32) |
| 42 | + num_y = np.random.randint(0, 10, 500).astype(dtype=np.int64) |
| 43 | + dataset = tf.data.Dataset.from_tensor_slices((num_x, num_y)) \ |
| 44 | + .batch(bs).repeat(1) |
| 45 | + iterator = dataset.make_initializable_iterator() |
| 46 | + tf.add_to_collection(tf.GraphKeys.TABLE_INITIALIZERS, iterator.initializer) |
| 47 | + x, labels = iterator.get_next() |
| 48 | + for _ in range(10): |
| 49 | + x = tf.layers.dense(x, 10, use_bias=False) |
| 50 | + logits = x |
| 51 | + loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits) |
| 52 | + if enable_epl: |
| 53 | + epl.add_to_collection(loss, epl.GraphKeys.GLOBAL_MEAN_OBJECTS) |
| 54 | + tf.summary.scalar('loss', loss) |
| 55 | + global_step = tf.train.get_or_create_global_step() |
| 56 | + learning_rate = 0.001 |
| 57 | + if lr_decay: |
| 58 | + learning_rate = tf.train.exponential_decay(learning_rate, global_step, |
| 59 | + max_steps, 0.96, |
| 60 | + staircase=False) |
| 61 | + if optimizer is None: |
| 62 | + optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate) |
| 63 | + else: |
| 64 | + optimizer = optimizer |
| 65 | + train_op = optimizer.minimize(loss, global_step=global_step) |
| 66 | + merged = tf.summary.merge_all() |
| 67 | + hooks = [tf.train.StopAtStepHook(last_step=max_steps)] |
| 68 | + losses = [] |
| 69 | + with tf.train.MonitoredTrainingSession( |
| 70 | + hooks=hooks) as sess: |
| 71 | + while not sess.should_stop(): |
| 72 | + fetches = [loss, train_op, global_step, merged] |
| 73 | + res = sess.run(fetches) |
| 74 | + train_loss = res[0] |
| 75 | + losses.append(train_loss) |
| 76 | + print("step: {} loss: {}".format(res[2], train_loss)) |
| 77 | + return losses, optimizer |
| 78 | + |
| 79 | +# pylint: disable=missing-docstring,unused-variable |
| 80 | +# pylint: disable=protected-access |
| 81 | +class MultiOptimizerTest(test.TestCase): |
| 82 | + @classmethod |
| 83 | + def setUpClass(cls): # pylint: disable=invalid-name |
| 84 | + with tf.Graph().as_default(): |
| 85 | + optimizer = get_nested_optimizer() |
| 86 | + cls.base, _ = _train_fn(bs=2, lr_decay=False, max_steps=5, optimizer=optimizer, enable_epl=False) |
| 87 | + |
| 88 | + |
| 89 | + def test_nested_optimizer(self): |
| 90 | + epl.init(epl.Config({"communication.clip_after_allreduce": True})) |
| 91 | + epl.set_default_strategy(epl.replicate(1)) |
| 92 | + with tf.Graph().as_default(): |
| 93 | + optimizer = get_nested_optimizer() |
| 94 | + res2, _ = _train_fn(bs=1, lr_decay=False, max_steps=5, optimizer=optimizer) |
| 95 | + self.assertEqual(len(epl.Graph.get().gradients), 10) |
| 96 | + |
| 97 | + def test_nested_optimizer_clip_after_allreduce(self): |
| 98 | + epl.init(epl.Config({"communication.clip_after_allreduce": False})) |
| 99 | + epl.set_default_strategy(epl.replicate(1)) |
| 100 | + with tf.Graph().as_default(): |
| 101 | + optimizer = get_nested_optimizer() |
| 102 | + res2, _ = _train_fn(bs=1, lr_decay=False, max_steps=5, optimizer=optimizer) |
| 103 | + self.assertEqual(len(epl.Graph.get().gradients), 10) |
| 104 | + |
| 105 | +# pylint: enable=missing-docstring,unused-variable |
| 106 | +# pylint: enable=protected-access |
| 107 | + |
| 108 | +if __name__ == "__main__": |
| 109 | + test.main() |
0 commit comments