def test_keras_gradients(script_mode, tf_optimizer):
""" Works as intended. """
smd.del_hook()
tf.reset_default_graph()
tf.keras.backend.clear_session()
json_file_contents = """
{
"S3OutputPath": "s3://sagemaker-test",
"LocalPath": "/opt/ml/output/tensors",
"CollectionConfigurations": [
{
"CollectionName": "gradients"
},
{
"CollectionName": "optimizer_variables"
},
{
"CollectionName": "losses"
}
]
}
"""
with SagemakerSimulator(json_file_contents=json_file_contents) as sim:
model = get_keras_model_v1()
(x_train, y_train), (x_test, y_test) = get_keras_data()
if tf_optimizer:
opt = tf.train.RMSPropOptimizer(0.1)
else:
opt = tf.keras.optimizers.RMSprop()
if script_mode:
hook = smd.KerasHook(
out_dir=sim.out_dir,
include_collections=["gradients", "optimizer_variables", "losses"],
)
opt = hook.wrap_optimizer(opt)
model.compile(
loss="sparse_categorical_crossentropy", optimizer=opt, metrics=["accuracy"]
)
history = model.fit(
x_train, y_train, batch_size=16, epochs=5, validation_split=0.2, callbacks=[hook]
)
test_scores = model.evaluate(x_test, y_test, verbose=2, callbacks=[hook])
else:
model.compile(
loss="sparse_categorical_crossentropy", optimizer=opt, metrics=["accuracy"]
)
history = model.fit(x_train, y_train, batch_size=16, epochs=5, validation_split=0.2)
test_scores = model.evaluate(x_test, y_test, verbose=2)
# Check that hook created and tensors saved
trial = smd.create_trial(path=sim.out_dir)
assert smd.get_hook() is not None, "Hook was not created."
assert len(trial.steps()) > 0, "Nothing saved at any step."
assert len(trial.tensor_names()) > 0, "Tensors were not saved."
assert len(trial.tensor_names(collection="gradients")) > 0
if not tf_optimizer:
# as this is only supported for keras optimizers currently
assert len(trial.tensor_names(collection="optimizer_variables")) > 0
def test_keras_v1(script_mode):
""" Works as intended. """
smd.del_hook()
tf.reset_default_graph()
tf.keras.backend.clear_session()
with SagemakerSimulator() as sim:
model = get_keras_model_v1()
(x_train, y_train), (x_test, y_test) = get_keras_data()
model.compile(
loss="sparse_categorical_crossentropy",
optimizer=tf.keras.optimizers.RMSprop(),
metrics=["accuracy"],
)
if script_mode:
hook = smd.KerasHook(out_dir=sim.out_dir)
history = model.fit(
x_train, y_train, batch_size=64, epochs=5, validation_split=0.2, callbacks=[hook]
)
test_scores = model.evaluate(x_test, y_test, verbose=2, callbacks=[hook])
else:
history = model.fit(x_train, y_train, batch_size=64, epochs=5, validation_split=0.2)
test_scores = model.evaluate(x_test, y_test, verbose=2)
# Check that hook created and tensors saved
trial = smd.create_trial(path=sim.out_dir)
assert smd.get_hook() is not None, "Hook was not created."
assert len(trial.steps()) > 0, "Nothing saved at any step."
assert len(trial.tensor_names()) > 0, "Tensors were not saved."