def compile_model(
context: keras.TFKerasContext,
compile_args: inspect.BoundArguments,
env: det.EnvContext,
hvd_config: horovod.HorovodContext,
) -> None:
(
context.model,
compile_args.arguments["optimizer"],
) = keras._get_multi_gpu_model_and_optimizer(
pre_compiled_model=context.model,
optimizer=compile_args.arguments["optimizer"],
env=env,
hvd_config=hvd_config,
profile_frequency=env.experiment_config.profile_frequency(),
profile_filename=DeterminedProfiler.OUTPUT_FILENAME,
)
if hvd_config.use and version.parse("2.0.0") <= version.parse(
tf.__version__) < version.parse("2.2.0"):
logging.info(
"Calling `model.compile(...)` with `experimental_run_tf_function=False` to ensure "
"TensorFlow calls `optimizer.get_gradients()` to compute gradients."
)
context.model.compile(*compile_args.args,
**compile_args.kwargs,
experimental_run_tf_function=False)
else:
context.model.compile(*compile_args.args, **compile_args.kwargs)
def from_native(
context: det.NativeContext,
env: det.EnvContext,
workloads: workload.Stream,
load_path: Optional[pathlib.Path],
rendezvous_info: det.RendezvousInfo,
hvd_config: horovod.HorovodContext,
) -> det.TrialController:
check.is_instance(
context,
keras.TFKerasNativeContext,
"TFKerasTrialController needs a TFKerasSprinkleContext",
)
context = cast(keras.TFKerasNativeContext, context)
check.is_not_none(context.model, "Please call wrap_model(...).")
check.is_not_none(context.compile_args,
"Please call model.compile(...).")
check.is_not_none(
context.train_config,
"Please call model.fit(...) or model.fit_generator(...).",
)
# For the Native API, we would break the user's model if we changed the session
# right now, so we have to trust the user did not modify what we set previously.
#
# TODO(ryan): Fix this, probably with a function for configuring the backend session.
session = tf.compat.v1.keras.backend.get_session()
compile_args = cast(inspect.BoundArguments, context.compile_args)
train_config = cast(keras.TFKerasTrainConfig, context.train_config)
(
context.model,
compile_args.arguments["optimizer"],
) = keras._get_multi_gpu_model_and_optimizer(
pre_compiled_model=context.model,
optimizer=compile_args.arguments["optimizer"],
env=env,
hvd_config=hvd_config,
profile_frequency=env.experiment_config.profile_frequency(),
profile_filename=DeterminedProfiler.OUTPUT_FILENAME,
)
context.model.compile(*compile_args.args, **compile_args.kwargs)
return TFKerasTrialController(
context.model,
session,
train_config,
context,
env,
workloads,
load_path,
rendezvous_info,
hvd_config,
)
def from_trial(
trial_inst: det.Trial,
context: det.TrialContext,
env: det.EnvContext,
workloads: workload.Stream,
load_path: Optional[pathlib.Path],
rendezvous_info: det.RendezvousInfo,
hvd_config: horovod.HorovodContext,
) -> det.TrialController:
check.is_instance(
context,
keras.TFKerasTrialContext,
"TFKerasTrialController needs a TFKerasTrialContext",
)
context = cast(keras.TFKerasTrialContext, context)
check.is_instance(trial_inst, TFKerasTrial,
"TFKerasTrialController needs a TFKerasTrial")
trial = cast(TFKerasTrial, trial_inst)
session = TFKerasTrialController._configure_session(
env, hvd_config, trial.session_config())
training_x, training_y, training_sample_weight = keras._get_x_y_and_sample_weight(
input_data=trial.build_training_data_loader())
training_data = keras._adapt_keras_data(
x=training_x,
y=training_y,
sample_weight=training_sample_weight,
batch_size=context.get_per_slot_batch_size(),
drop_leftovers=True,
)
val_x, val_y, val_sample_weight = keras._get_x_y_and_sample_weight(
input_data=trial.build_validation_data_loader())
validation_data = keras._adapt_keras_data(
x=val_x,
y=val_y,
sample_weight=val_sample_weight,
batch_size=context.get_per_slot_batch_size(),
drop_leftovers=False,
)
trial.build_model()
check.is_not_none(context.model, "Please call wrap_model(...).")
check.is_not_none(context.compile_args,
"Please call model.compile(...).")
compile_args = cast(inspect.BoundArguments, context.compile_args)
(
context.model,
compile_args.arguments["optimizer"],
) = keras._get_multi_gpu_model_and_optimizer(
pre_compiled_model=context.model,
optimizer=compile_args.arguments["optimizer"],
env=env,
hvd_config=hvd_config,
profile_frequency=env.experiment_config.profile_frequency(),
profile_filename=DeterminedProfiler.OUTPUT_FILENAME,
)
if hvd_config.use and version.parse(
tf.__version__) >= version.parse("2.0.0"):
logging.info(
"Calling `model.compile(...)` with `experimental_run_tf_function=False` to ensure "
"TensorFlow calls `optimizer.get_gradients()` to compute gradients."
)
context.model.compile(*compile_args.args,
**compile_args.kwargs,
experimental_run_tf_function=False)
else:
context.model.compile(*compile_args.args, **compile_args.kwargs)
tf_keras_callbacks = trial.keras_callbacks()
return TFKerasTrialController(
context.model,
session,
keras.TFKerasTrainConfig(training_data, validation_data,
tf_keras_callbacks),
context,
env,
workloads,
load_path,
rendezvous_info,
hvd_config,
)
