def testMovingAverage(self): step = tf.Variable(0, trainable=False, dtype=tf.int64) variables = [tf.Variable(1.0), tf.Variable(2.0)] moving_average = training.MovingAverage(variables, step) variables[0].assign(3.0) variables[1].assign(4.0) moving_average.update() self.assertAllEqual(self.evaluate(variables), [3.0, 4.0]) with moving_average.shadow_variables(): self.assertAllClose(self.evaluate(variables), [2.8, 3.8]) self.assertAllEqual(self.evaluate(variables), [3.0, 4.0])
def testMovingAverageDistributionStrategy(self): devices = tf.config.experimental.list_logical_devices(device_type="CPU") strategy = tf.distribute.MirroredStrategy(devices=devices) with strategy.scope(): variables = [tf.Variable(1.0), tf.Variable(2.0)] step = tf.Variable(0, trainable=False, dtype=tf.int64) moving_average = training.MovingAverage(variables, step) variables[0].assign(3.0) variables[1].assign(4.0) moving_average.update() self.assertAllEqual(self.evaluate(variables), [3.0, 4.0]) with moving_average.shadow_variables(): self.assertAllClose(self.evaluate(variables), [2.8, 3.8]) self.assertAllEqual(self.evaluate(variables), [3.0, 4.0])