def test_train_sequential_with_distribution_strategy(self):
    dist = mirrored_strategy.MirroredStrategy(devices=['/device:GPU:0', '/device:GPU:1'])
    keras_model = simple_sequential_model()
    keras_model.compile(loss='categorical_crossentropy', metrics=[keras.metrics.CategoricalAccuracy()], optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.01))
    config = run_config_lib.RunConfig(tf_random_seed=_RANDOM_SEED, model_dir=self._base_dir, train_distribute=dist)
    with self.cached_session():
        est_keras = keras_lib.model_to_estimator(keras_model=keras_model, config=config)
        before_eval_results = est_keras.evaluate(input_fn=get_ds_test_input_fn, steps=1)
        est_keras.train(input_fn=get_ds_train_input_fn, steps=(_TRAIN_SIZE / 16))
        after_eval_results = est_keras.evaluate(input_fn=get_ds_test_input_fn, steps=1)
        self.assertLess(after_eval_results['loss'], before_eval_results['loss'])
    writer_cache.FileWriterCache.clear()
    gfile.DeleteRecursively(self._config.model_dir)
def test_train_sequential_with_distribution_strategy(self):
    dist = mirrored_strategy.MirroredStrategy(devices=['/device:GPU:0', '/device:GPU:1'])
    keras_model = simple_sequential_model()
    keras_model.compile(loss='categorical_crossentropy', metrics=[keras.metrics.CategoricalAccuracy()], optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.01))
    config = run_config_lib.RunConfig(tf_random_seed=_RANDOM_SEED, model_dir=self._base_dir, train_distribute=dist)
    with self.cached_session():
        est_keras = keras_lib.model_to_estimator(keras_model=keras_model, config=config)
        before_eval_results = est_keras.evaluate(input_fn=get_ds_test_input_fn, steps=1)
        est_keras.train(input_fn=get_ds_train_input_fn, steps=(_TRAIN_SIZE / 16))
        after_eval_results = est_keras.evaluate(input_fn=get_ds_test_input_fn, steps=1)
        self.assertLess(after_eval_results['loss'], before_eval_results['loss'])
    writer_cache.FileWriterCache.clear()
    gfile.DeleteRecursively(self._config.model_dir)