'Given keras model has not been compiled yet. Please compile first '
'before creating the estimator.')
- keras_weights = keras_model.get_weights()
keras_model_fn = _create_keras_model_fn(keras_model, custom_objects)
est = estimator_lib.Estimator(
keras_model_fn, model_dir=model_dir, config=config)
+ # Pass the config into keras backend's default session.
+ with session.Session(config=est._session_config) as sess:
+ K.set_session(sess)
+
+ keras_weights = keras_model.get_weights()
# TODO(yifeif): move checkpoint initialization to scaffold.init_fn
_save_first_checkpoint(keras_model, est, custom_objects, keras_weights)
return est
import numpy as np
+from tensorflow.core.protobuf import config_pb2
from tensorflow.python.estimator import run_config as run_config_lib
from tensorflow.python.estimator.inputs import numpy_io
from tensorflow.python.framework import test_util
keras_model=keras_model,
model_dir=tempfile.mkdtemp(dir=self._base_dir))
+ def test_gpu_config(self):
+ keras_model, (_, _), (_, _), _, _ = get_resource_for_simple_model()
+ keras_model.compile(
+ loss='categorical_crossentropy',
+ optimizer='rmsprop',
+ metrics=['mse', keras.metrics.categorical_accuracy])
+
+ gpu_options = config_pb2.GPUOptions(per_process_gpu_memory_fraction=0.3)
+ sess_config = config_pb2.ConfigProto(gpu_options=gpu_options)
+ self._config._session_config = sess_config
+ keras.estimator.model_to_estimator(
+ keras_model=keras_model, config=self._config)
+ self.assertEqual(keras.backend.get_session()
+ ._config.gpu_options.per_process_gpu_memory_fraction,
+ gpu_options.per_process_gpu_memory_fraction)
+
if __name__ == '__main__':
test.main()