def main(_):
# Build the train and eval datasets from the MNIST data. Also return the
# input shape which is constructed based on the `image_data_format`
# i.e channels_first or channels_last.
tf.enable_eager_execution()
train_ds, eval_ds, input_shape = get_input_datasets()
# Instantiate the MirroredStrategy object. If we don't specify `num_gpus` or
# the `devices` argument then all the GPUs available on the machine are used.
# TODO(priyag): Use `tf.distribute.MirroredStrategy` once available.
strategy = mirrored_strategy.MirroredStrategy(['/gpu:0', '/cpu:0'])
# Create and compile the model under Distribution strategy scope.
# `fit`, `evaluate` and `predict` will be distributed based on the strategy
# model was compiled with.
with strategy.scope():
model = get_model(input_shape)
optimizer = rmsprop.RMSProp(learning_rate=0.001)
model.compile(loss=tf.keras.losses.categorical_crossentropy,
optimizer=optimizer,
metrics=['accuracy'])
# Train the model with the train dataset.
model.fit(x=train_ds, epochs=20, steps_per_epoch=468)
# Evaluate the model with the eval dataset.
score = model.evaluate(eval_ds, steps=10, verbose=0)
print('Test loss:', score[0])
print('Test accuracy:', score[1])
def fit_with(input_shape, verbose, dropout2_rate, dense_1_neurons_x128, lr):
# Create the model using a specified hyperparameters.
dense_1_neurons = max(int(dense_1_neurons_x128 * 128), 128)
model = get_model(input_shape, dropout2_rate, dense_1_neurons)
# Train the model for a specified number of epochs.
optimizer = rmsprop.RMSProp(learning_rate=lr)
model.compile(loss=tf.keras.losses.categorical_crossentropy,
optimizer=optimizer,
metrics=['accuracy'])
# Train the model with the train dataset.
model.fit(x=train_ds,
epochs=1,
steps_per_epoch=468,
batch_size=64,
verbose=verbose)
# Evaluate the model with the eval dataset.
score = model.evaluate(eval_ds, steps=10, verbose=0)
print('Test loss:', score[0])
print('Test accuracy:', score[1])
# Return the accuracy.
return score[1]
def get_model_saveable(implementation, filters, kernel_size, strides, layers,
num_classes, data_format):
model = keras.Sequential()
if len(kernel_size) == 1:
lc_layer = keras.layers.LocallyConnected1D
elif len(kernel_size) == 2:
lc_layer = keras.layers.LocallyConnected2D
else:
raise NotImplementedError(kernel_size)
for _ in range(layers):
model.add(
lc_layer(padding='valid',
kernel_initializer=keras.initializers.random_normal(),
bias_initializer=keras.initializers.random_normal(),
filters=filters,
strides=strides,
kernel_size=kernel_size,
activation=keras.activations.relu,
data_format=data_format,
implementation=implementation))
model.add(keras.layers.Flatten())
model.add(keras.layers.Dense(num_classes))
model.compile(optimizer=rmsprop.RMSProp(learning_rate=0.01),
metrics=[keras.metrics.categorical_accuracy],
loss=xent)
return model
def testCallableParams(self):
with context.eager_mode():
for dtype in [dtypes.half, dtypes.float32]:
var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype)
grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
learning_rate = lambda: 2.0
rho = lambda: 0.9
momentum = lambda: 0.0
epsilon = lambda: 1.0
opt = rmsprop.RMSProp(learning_rate, rho, momentum, epsilon)
# Fetch params to validate initial values
self.assertAllClose([1.0, 2.0], self.evaluate(var0))
self.assertAllClose([3.0, 4.0], self.evaluate(var1))
# Step 1: the rms accumulators where 1. So we should see a normal
# update: v -= grad * learning_rate
opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
# Check the parameters.
self.assertAllCloseAccordingToType(
np.array([
1.0 - (0.1 * 2.0 / math.sqrt(0.001 + 1.0)),
2.0 - (0.1 * 2.0 / math.sqrt(0.001 + 1.0))
]), self.evaluate(var0))
self.assertAllCloseAccordingToType(
np.array([
3.0 - (0.01 * 2.0 / math.sqrt(0.00001 + 1.0)),
4.0 - (0.01 * 2.0 / math.sqrt(0.00001 + 1.0))
]), self.evaluate(var1))
# Step 2: the root mean square accumulators contain the previous update.
opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
# Check the parameters.
self.assertAllCloseAccordingToType(
np.array([
1.0 - (0.1 * 2.0 / math.sqrt(0.001 + 1.0)) -
(0.1 * 2.0 / math.sqrt(0.001 * 0.9 + 0.001 + 1.0)),
2.0 - (0.1 * 2.0 / math.sqrt(0.001 + 1.0)) -
(0.1 * 2.0 / math.sqrt(0.001 * 0.9 + 0.001 + 1.0))
]), self.evaluate(var0))
self.assertAllCloseAccordingToType(
np.array([
3.0 - (0.01 * 2.0 / math.sqrt(0.00001 + 1.0)) -
(0.01 * 2.0 / math.sqrt(0.00001 * 0.9 + 1e-5 + 1.0)),
4.0 - (0.01 * 2.0 / math.sqrt(0.00001 + 1.0)) -
(0.01 * 2.0 / math.sqrt(0.00001 * 0.9 + 1e-5 + 1.0))
]), self.evaluate(var1))
def testConfig(self):
def momentum():
return ops.convert_to_tensor(3.0)
opt = rmsprop.RMSProp(learning_rate=1.0,
rho=2.0,
momentum=momentum,
epsilon=lambda: ops.convert_to_tensor(4.0),
centered=True)
config = opt.get_config()
opt2 = rmsprop.RMSProp.from_config(config)
self.assertEqual(opt._hyper["learning_rate"][1],
opt2._hyper["learning_rate"][1])
self.assertEqual(opt._hyper["rho"][1], opt2._hyper["rho"][1])
self.assertEqual(opt._hyper["momentum"][1].__name__,
opt2._hyper["momentum"][1].__name__)
self.assertIsInstance(opt2._hyper["epsilon"][1],
python_types.LambdaType)
self.assertEqual(True, opt2._centered)
def main(_):
if flags.FLAGS.enable_eager:
ops.enable_eager_execution()
logging.info('Eager execution enabled for MNIST Multi-Worker.')
else:
logging.info('Eager execution not enabled for MNIST Multi-Worker.')
# Build the train and eval datasets from the MNIST data.
train_ds, eval_ds = get_input_datasets()
if flags.FLAGS.distribution_strategy == 'multi_worker_mirrored':
# MultiWorkerMirroredStrategy for multi-worker distributed MNIST training.
strategy = collective_strategy.CollectiveAllReduceStrategy()
else:
raise ValueError(
'Only `multi_worker_mirrored` is supported strategy '
'in Keras MNIST example at this time. Strategy passed '
'in is %s' % flags.FLAGS.distribution_strategy)
# Create and compile the model under Distribution strategy scope.
# `fit`, `evaluate` and `predict` will be distributed based on the strategy
# model was compiled with.
with strategy.scope():
model = get_model()
optimizer = rmsprop.RMSProp(learning_rate=0.001)
model.compile(loss=keras.losses.categorical_crossentropy,
optimizer=optimizer,
metrics=['accuracy'])
# Train the model with the train dataset.
tensorboard_callback = keras.callbacks.TensorBoard(
log_dir=flags.FLAGS.model_dir)
model.fit(x=train_ds,
epochs=20,
steps_per_epoch=468,
callbacks=[tensorboard_callback])
# Evaluate the model with the eval dataset.
score = model.evaluate(eval_ds, steps=10, verbose=0)
logging.info('Test loss:{}'.format(score[0]))
logging.info('Test accuracy:{}'.format(score[1]))
def testMinimizeSparseResourceVariableCentered(self):
for dtype in [dtypes.float32, dtypes.float64]:
with self.cached_session():
var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype)
x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
loss = pred * pred
sgd_op = rmsprop.RMSProp(
learning_rate=1.0,
rho=0.0,
momentum=0.0,
epsilon=1.0,
centered=True).minimize(
loss, var_list=[var0])
variables.global_variables_initializer().run()
# Fetch params to validate initial values
self.assertAllCloseAccordingToType([[1.0, 2.0]], var0.eval())
# Run 1 step of sgd
sgd_op.run()
# Validate updated params
self.assertAllCloseAccordingToType(
[[-111, -138]], var0.eval(), atol=0.01)
def testDense(self):
for (dtype, learning_rate, rho, momentum, epsilon, centered) in _TESTPARAMS:
with self.cached_session(use_gpu=True):
# Initialize variables for numpy implementation.
var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
grads0_np = np.array([0.1, 0.2], dtype=dtype.as_numpy_dtype)
var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
grads1_np = np.array([0.01, 0.2], dtype=dtype.as_numpy_dtype)
var0 = resource_variable_ops.ResourceVariable(var0_np, dtype=dtype)
var1 = resource_variable_ops.ResourceVariable(var1_np, dtype=dtype)
grads0 = constant_op.constant(grads0_np, dtype=dtype)
grads1 = constant_op.constant(grads1_np, dtype=dtype)
opt = rmsprop.RMSProp(
learning_rate=learning_rate,
rho=rho,
momentum=momentum,
epsilon=epsilon,
centered=centered)
update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
variables.global_variables_initializer().run()
if centered:
mg0 = opt.get_slot(var0, "mg")
mg1 = opt.get_slot(var1, "mg")
else:
mg0 = None
mg1 = None
rms0 = opt.get_slot(var0, "rms")
self.assertTrue(rms0 is not None)
rms1 = opt.get_slot(var1, "rms")
self.assertTrue(rms1 is not None)
mom0 = opt.get_slot(var0, "momentum")
self.assertTrue(mom0 is not None)
mom1 = opt.get_slot(var1, "momentum")
self.assertTrue(mom1 is not None)
mg0_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
mg1_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
rms0_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
rms1_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
mom0_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
mom1_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
# Fetch params to validate initial values
self.assertAllClose([1.0, 2.0], var0.eval())
self.assertAllClose([3.0, 4.0], var1.eval())
# Run 4 steps of RMSProp
for _ in range(1, 5):
update.run()
var0_np, mg0_np, rms0_np, mom0_np = self._rmsprop_update_numpy(
var0_np, grads0_np, mg0_np, rms0_np, mom0_np, learning_rate, rho,
momentum, epsilon, centered)
var1_np, mg1_np, rms1_np, mom1_np = self._rmsprop_update_numpy(
var1_np, grads1_np, mg1_np, rms1_np, mom1_np, learning_rate, rho,
momentum, epsilon, centered)
# Validate updated params
if centered:
self.assertAllCloseAccordingToType(mg0_np, mg0.eval())
self.assertAllCloseAccordingToType(mg1_np, mg1.eval())
self.assertAllCloseAccordingToType(rms0_np, rms0.eval())
self.assertAllCloseAccordingToType(rms1_np, rms1.eval())
self.assertAllCloseAccordingToType(mom0_np, mom0.eval())
self.assertAllCloseAccordingToType(mom1_np, mom1.eval())
self.assertAllCloseAccordingToType(var0_np, var0.eval())
self.assertAllCloseAccordingToType(var1_np, var1.eval())
def testDense(self, dtype, param_value):
(learning_rate, rho, momentum, epsilon, centered,
use_resource) = tuple(param_value)
with self.test_session(use_gpu=True):
# Initialize variables for numpy implementation.
var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
grads0_np = np.array([0.1, 0.2], dtype=dtype.as_numpy_dtype)
var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
grads1_np = np.array([0.01, 0.2], dtype=dtype.as_numpy_dtype)
if use_resource:
var0 = resource_variable_ops.ResourceVariable(var0_np)
var1 = resource_variable_ops.ResourceVariable(var1_np)
else:
var0 = variables.Variable(var0_np)
var1 = variables.Variable(var1_np)
grads0 = constant_op.constant(grads0_np)
grads1 = constant_op.constant(grads1_np)
opt = rmsprop.RMSProp(learning_rate=learning_rate,
rho=rho,
momentum=momentum,
epsilon=epsilon,
centered=centered)
update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
variables.global_variables_initializer().run()
mg0 = opt.get_slot(var0, "mg")
self.assertEqual(mg0 is not None, centered)
mg1 = opt.get_slot(var1, "mg")
self.assertEqual(mg1 is not None, centered)
rms0 = opt.get_slot(var0, "rms")
self.assertIsNotNone(rms0)
rms1 = opt.get_slot(var1, "rms")
self.assertIsNotNone(rms1)
mom0 = opt.get_slot(var0, "momentum")
self.assertIsNotNone(mom0)
mom1 = opt.get_slot(var1, "momentum")
self.assertIsNotNone(mom1)
mg0_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
mg1_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
rms0_np = np.array([epsilon, epsilon], dtype=dtype.as_numpy_dtype)
rms1_np = np.array([epsilon, epsilon], dtype=dtype.as_numpy_dtype)
mom0_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
mom1_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
# Fetch params to validate initial values
self.assertAllClose([1.0, 2.0], var0.eval())
self.assertAllClose([3.0, 4.0], var1.eval())
# Run 4 steps of RMSProp
for _ in range(4):
update.run()
var0_np, mg0_np, rms0_np, mom0_np = self._rmsprop_update_numpy(
var0_np, grads0_np, mg0_np, rms0_np, mom0_np,
learning_rate, rho, momentum, centered)
var1_np, mg1_np, rms1_np, mom1_np = self._rmsprop_update_numpy(
var1_np, grads1_np, mg1_np, rms1_np, mom1_np,
learning_rate, rho, momentum, centered)
# Validate updated params
if centered:
self.assertAllCloseAccordingToType(mg0_np, mg0.eval())
self.assertAllCloseAccordingToType(mg1_np, mg1.eval())
self.assertAllCloseAccordingToType(rms0_np, rms0.eval())
self.assertAllCloseAccordingToType(rms1_np, rms1.eval())
self.assertAllCloseAccordingToType(mom0_np, mom0.eval())
self.assertAllCloseAccordingToType(mom1_np, mom1.eval())
self.assertAllCloseAccordingToType(var0_np,
var0.eval(),
half_rtol=0.01,
half_atol=0.01)
self.assertAllCloseAccordingToType(var1_np,
var1.eval(),
half_rtol=0.01,
half_atol=0.01)
def testWithMomentum(self, dtype):
with self.test_session(use_gpu=True):
var0 = variables.Variable([1.0, 2.0], dtype=dtype)
var1 = variables.Variable([3.0, 4.0], dtype=dtype)
grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
opt = rmsprop.RMSProp(learning_rate=2.0,
rho=0.9,
momentum=0.5,
epsilon=1.0)
update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
variables.global_variables_initializer().run()
rms0 = opt.get_slot(var0, "rms")
self.assertIsNotNone(rms0)
rms1 = opt.get_slot(var1, "rms")
self.assertIsNotNone(rms1)
mom0 = opt.get_slot(var0, "momentum")
self.assertIsNotNone(mom0)
mom1 = opt.get_slot(var1, "momentum")
self.assertIsNotNone(mom1)
# Fetch params to validate initial values
self.assertAllClose([1.0, 2.0], var0.eval())
self.assertAllClose([3.0, 4.0], var1.eval())
# Step 1: rms = 1, mom = 0. So we should see a normal
# update: v -= grad * learning_rate
update.run()
# Check the root mean square accumulators.
self.assertAllCloseAccordingToType(np.array([0.901, 0.901]),
rms0.eval())
self.assertAllCloseAccordingToType(np.array([0.90001, 0.90001]),
rms1.eval())
# Check the momentum accumulators
self.assertAllCloseAccordingToType(
np.array([(0.1 * 2.0 / math.sqrt(0.901)),
(0.1 * 2.0 / math.sqrt(0.901))]), mom0.eval())
self.assertAllCloseAccordingToType(
np.array([(0.01 * 2.0 / math.sqrt(0.90001)),
(0.01 * 2.0 / math.sqrt(0.90001))]), mom1.eval())
# Check that the parameters.
self.assertAllCloseAccordingToType(
np.array([
1.0 - (0.1 * 2.0 / math.sqrt(0.901)),
2.0 - (0.1 * 2.0 / math.sqrt(0.901))
]), var0.eval())
self.assertAllCloseAccordingToType(
np.array([
3.0 - (0.01 * 2.0 / math.sqrt(0.90001)),
4.0 - (0.01 * 2.0 / math.sqrt(0.90001))
]), var1.eval())
# Step 2: the root mean square accumulators contain the previous update.
update.run()
# Check the rms accumulators.
self.assertAllCloseAccordingToType(
np.array([0.901 * 0.9 + 0.001, 0.901 * 0.9 + 0.001]),
rms0.eval())
self.assertAllCloseAccordingToType(
np.array([0.90001 * 0.9 + 1e-5, 0.90001 * 0.9 + 1e-5]),
rms1.eval())
self.assertAllCloseAccordingToType(
np.array([
0.5 * (0.1 * 2.0 / math.sqrt(0.901)) +
(0.1 * 2.0 / math.sqrt(0.901 * 0.9 + 0.001)),
0.5 * (0.1 * 2.0 / math.sqrt(0.901)) +
(0.1 * 2.0 / math.sqrt(0.901 * 0.9 + 0.001))
]), mom0.eval())
self.assertAllCloseAccordingToType(
np.array([
0.5 * (0.01 * 2.0 / math.sqrt(0.90001)) +
(0.01 * 2.0 / math.sqrt(0.90001 * 0.9 + 1e-5)),
0.5 * (0.01 * 2.0 / math.sqrt(0.90001)) +
(0.01 * 2.0 / math.sqrt(0.90001 * 0.9 + 1e-5))
]), mom1.eval())
# Check the parameters.
self.assertAllCloseAccordingToType(
np.array([
1.0 - (0.1 * 2.0 / math.sqrt(0.901)) -
(0.5 * (0.1 * 2.0 / math.sqrt(0.901)) +
(0.1 * 2.0 / math.sqrt(0.901 * 0.9 + 0.001))),
2.0 - (0.1 * 2.0 / math.sqrt(0.901)) -
(0.5 * (0.1 * 2.0 / math.sqrt(0.901)) +
(0.1 * 2.0 / math.sqrt(0.901 * 0.9 + 0.001)))
]), var0.eval())
self.assertAllCloseAccordingToType(
np.array([
3.0 - (0.01 * 2.0 / math.sqrt(0.90001)) -
(0.5 * (0.01 * 2.0 / math.sqrt(0.90001)) +
(0.01 * 2.0 / math.sqrt(0.90001 * 0.9 + 1e-5))),
4.0 - (0.01 * 2.0 / math.sqrt(0.90001)) -
(0.5 * (0.01 * 2.0 / math.sqrt(0.90001)) +
(0.01 * 2.0 / math.sqrt(0.90001 * 0.9 + 1e-5)))
]), var1.eval())
def testSparse(self):
for (dtype, learning_rate, rho, momentum, epsilon, centered) in _TESTPARAMS:
with self.cached_session(use_gpu=True):
# Initialize variables for numpy implementation.
var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
grads0_np = np.array([0.1], dtype=dtype.as_numpy_dtype)
var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
grads1_np = np.array([0.01], dtype=dtype.as_numpy_dtype)
var0 = variables.Variable(var0_np)
var1 = variables.Variable(var1_np)
grads0_np_indices = np.array([0], dtype=np.int32)
grads0 = ops.IndexedSlices(
constant_op.constant(grads0_np),
constant_op.constant(grads0_np_indices), constant_op.constant([1]))
grads1_np_indices = np.array([1], dtype=np.int32)
grads1 = ops.IndexedSlices(
constant_op.constant(grads1_np),
constant_op.constant(grads1_np_indices), constant_op.constant([1]))
opt = rmsprop.RMSProp(
learning_rate=learning_rate,
rho=rho,
momentum=momentum,
epsilon=epsilon,
centered=centered)
update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
variables.global_variables_initializer().run()
if centered:
mg0 = opt.get_slot(var0, "mg")
self.assertEqual(mg0 is not None, centered)
mg1 = opt.get_slot(var1, "mg")
self.assertEqual(mg1 is not None, centered)
else:
mg0 = None
mg1 = None
rms0 = opt.get_slot(var0, "rms")
self.assertTrue(rms0 is not None)
rms1 = opt.get_slot(var1, "rms")
self.assertTrue(rms1 is not None)
mom0 = opt.get_slot(var0, "momentum")
self.assertTrue(mom0 is not None)
mom1 = opt.get_slot(var1, "momentum")
self.assertTrue(mom1 is not None)
mg0_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
mg1_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
rms0_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
rms1_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
mom0_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
mom1_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
# Fetch params to validate initial values
self.assertAllClose([1.0, 2.0], self.evaluate(var0))
self.assertAllClose([3.0, 4.0], self.evaluate(var1))
# Run 4 steps of RMSProp
for _ in range(1, 5):
update.run()
var0_np, mg0_np, rms0_np, mom0_np = self._sparse_rmsprop_update_numpy(
var0_np, grads0_np_indices, grads0_np, mg0_np, rms0_np, mom0_np,
learning_rate, rho, momentum, epsilon, centered)
var1_np, mg1_np, rms1_np, mom1_np = self._sparse_rmsprop_update_numpy(
var1_np, grads1_np_indices, grads1_np, mg1_np, rms1_np, mom1_np,
learning_rate, rho, momentum, epsilon, centered)
# Validate updated params
if centered:
self.assertAllCloseAccordingToType(mg0_np, self.evaluate(mg0))
self.assertAllCloseAccordingToType(mg1_np, self.evaluate(mg1))
self.assertAllCloseAccordingToType(rms0_np, self.evaluate(rms0))
self.assertAllCloseAccordingToType(rms1_np, self.evaluate(rms1))
self.assertAllCloseAccordingToType(mom0_np, self.evaluate(mom0))
self.assertAllCloseAccordingToType(mom1_np, self.evaluate(mom1))
self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
