def test_apply_gradients(self):
x = variable_scope.get_variable("x", initializer=1., dtype=dtypes.float32)
dataset = dataset_ops.Dataset.from_tensor_slices([np.nan, np.inf, 0.1])
itr = dataset.make_one_shot_iterator()
lr = 1
opt = gd.GradientDescentOptimizer(lr)
lsm = lsm_lib.FixedLossScaleManager(1.e4)
opt = lso.LossScaleOptimizer(opt, lsm)
train_fn = lambda: opt.apply_gradients([(itr.get_next(), x)])
if not context.executing_eagerly():
train_op = train_fn()
expected_output = [1, 1, 1 - 0.1]
actual_output = []
self.evaluate(variables.global_variables_initializer())
for _ in range(3):
# nan or inf is not applied.
if context.executing_eagerly():
train_fn()
else:
self.evaluate(train_op)
actual_output.append(self.evaluate(x))
self.assertAllClose(expected_output, actual_output)
def test_basic(self):
itr = _GetExampleIter([True] * 10 + [False] * 10)
loss_scale = 1000
lsm = lsm_lib.FixedLossScaleManager(loss_scale)
update_fn = lambda: lsm.update_loss_scale(itr.get_next())
self.evaluate(variables.global_variables_initializer())
if not context.executing_eagerly():
update_op = update_fn()
for _ in range(10):
if context.executing_eagerly():
update_fn()
else:
self.evaluate(update_op)
self.assertEqual(loss_scale, self.evaluate(lsm.get_loss_scale()))
def loss_scale_opt_fn(opt): return lso.LossScaleOptimizer(opt, lsm_lib.FixedLossScaleManager(1e4))