def helper_test_keras_v2_gradienttape(script_mode: bool = False,
json_file_contents="{}",
default=False):
""" Test the default ZCC behavior of saving losses and metrics in eager and non-eager modes."""
smd.del_hook()
tf.keras.backend.clear_session()
with SagemakerSimulator(json_file_contents=json_file_contents) as sim:
helper_keras_gradienttape_train(script_mode=script_mode,
json_file_contents=json_file_contents,
sim=sim)
hook = smd.get_hook()
if script_mode:
assert hook
if default:
assert hook.has_default_hook_configuration()
hook.close()
# Check that hook created and tensors saved
trial = smd.create_trial(path=sim.out_dir)
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="losses")) > 0
else:
if version.parse(tf.__version__) < version.parse("2.1.2"):
assert not hook # only supported on TF 2.1.2 and greater
return
assert hook
hook.close()
# Check that hook created and tensors saved
trial = smd.create_trial(path=sim.out_dir)
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="losses")) > 0
if is_tf_2_2() and default is False:
# Inputs and Outputs are not saved with the default collection configurations.
assert len(trial.tensor_names(collection="inputs")) > 0
assert len(trial.tensor_names(collection="outputs")) > 0
assert trial.tensor_names(collection="outputs") == [
"predictions"
]
if "dense_layers" in json_file_contents:
# Only assert for test_keras_v2_multi_collections
# which defines this custom collection
assert len(
trial.tensor_names(collection="dense_layers")) > 0
else:
assert len(
trial.tensor_names(collection="dense_layers")) == 0
def test_keras_fit(out_dir, tf_eager_mode, saveall):
hook = smd.KerasHook(out_dir=out_dir, save_all=saveall)
helper_keras_fit(
trial_dir=out_dir,
hook=hook,
eager=tf_eager_mode,
steps=["train", "eval", "predict", "train"],
)
trial = smd.create_trial(path=out_dir)
# can't save gradients in TF 2.x eager mode
if saveall: # save losses, metrics, weights, biases
if tf_eager_mode:
assert len(trial.tensor_names()) == (12 if is_tf_2_2() else 13)
else:
assert len(trial.tensor_names()) == 21
assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 2
assert len(
trial.tensor_names(
collection=CollectionKeys.OPTIMIZER_VARIABLES)) == 5
assert (
len(
trial.tensor_names(
collection=CollectionKeys.OPTIMIZER_VARIABLES,
mode=ModeKeys.EVAL)) == 0,
"No Optimizer Variables Should be Saved in EVAL Mode",
)
else: # save the default losses and metrics
assert len(trial.tensor_names()) == (3 if is_tf_2_2() and tf_eager_mode
else 4)
assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 1
assert len(trial.tensor_names(collection=CollectionKeys.METRICS)) == (
2 if is_tf_2_2() and tf_eager_mode else 3)
def test_subclassed_model(out_dir):
# Download and load MNIST dataset.
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data("MNIST-data")
x_train, x_test = x_train / 255.0, x_test / 255.0
# Add a channels dimension
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]
# Create an instance of the model
model = MyModel()
train_ds = (
tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(10000, seed=123).batch(2)
)
MyModel.hook = smd.KerasHook(
out_dir,
save_all=True,
save_config=smd.SaveConfig(save_steps=[x for x in range(10)], save_interval=1),
)
MyModel.hook.register_model(model)
model.compile(optimizer="Adam", loss="mse", run_eagerly=True)
model.fit(train_ds, epochs=1, steps_per_epoch=10, callbacks=[MyModel.hook])
trial = smd.create_trial(out_dir)
assert len(trial.tensor_names(collection=smd.CollectionKeys.LAYERS)) == 8
assert trial.tensor_names(collection=smd.CollectionKeys.INPUTS) == ["model_input"]
assert trial.tensor_names(collection=smd.CollectionKeys.OUTPUTS) == ["labels", "predictions"]
assert trial.tensor_names(collection=smd.CollectionKeys.LOSSES) == ["loss"]
assert len(trial.tensor_names(collection=smd.CollectionKeys.GRADIENTS)) == 6
def test_estimator(out_dir, tf_eager_mode, saveall):
""" Works as intended. """
if tf_eager_mode is False:
tf.compat.v1.disable_eager_execution()
tf.compat.v1.reset_default_graph()
tf.keras.backend.clear_session()
mnist_classifier = get_estimator()
train_input_fn, eval_input_fn = get_input_fns()
# Train and evaluate
train_steps, eval_steps = 8, 2
hook = smd.EstimatorHook(out_dir=out_dir, save_all=saveall)
hook.set_mode(mode=smd.modes.TRAIN)
mnist_classifier.train(input_fn=train_input_fn, steps=train_steps, hooks=[hook])
hook.set_mode(mode=smd.modes.EVAL)
mnist_classifier.evaluate(input_fn=eval_input_fn, steps=eval_steps, hooks=[hook])
# Check that hook created and tensors saved
trial = smd.create_trial(path=out_dir)
tnames = trial.tensor_names()
assert len(trial.steps()) > 0
if saveall:
# Number of tensors in each collection
# vanilla TF 2.2: all = 300, loss = 1, weights = 4, gradients = 0, biases = 18, optimizer variables = 0, metrics = 0, others = 277
# AWS-TF 2.2 : all = 300, loss = 1, weights = 4, gradients = 8, biases = 18, optimizer variables = 0, metrics = 0, others = 269
# AWS-TF 2.1 : all = 309, loss = 1, weights = 4, gradients = 8, biases = 18, optimizer variables = 0, metrics = 0, others = 278
assert len(tnames) >= 300
assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 1
assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 4
assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 18
assert len(trial.tensor_names(collection=CollectionKeys.GRADIENTS)) >= 0
assert len(trial.tensor_names(collection=CollectionKeys.OPTIMIZER_VARIABLES)) >= 0
else:
assert len(tnames) == 1
assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 1
def test_linear_classifier(out_dir, tf_eager_mode, saveall):
""" Works as intended. """
if tf_eager_mode is False:
tf.compat.v1.disable_eager_execution()
tf.compat.v1.reset_default_graph()
tf.keras.backend.clear_session()
train_input_fn, eval_input_fn = get_input_fns()
x_feature = tf.feature_column.numeric_column("x", shape=(28, 28))
estimator = tf.estimator.LinearClassifier(
feature_columns=[x_feature], model_dir="/tmp/mnist_linear_classifier", n_classes=10
)
hook = smd.EstimatorHook(out_dir=out_dir, save_all=saveall)
estimator.train(input_fn=train_input_fn, steps=10, hooks=[hook])
# Check that hook created and tensors saved
trial = smd.create_trial(path=out_dir)
tnames = trial.tensor_names()
assert len(trial.steps()) > 0
if saveall:
# Number of tensors in each collection
# vanilla TF 2.2: all = 214, loss = 2, weights = 1, gradients = 0, biases = 12, optimizer variables = 0, metrics = 0, others = 199
# AWS-TF 2.2: all = 219, loss = 2, weights = 1, gradients = 2, biases = 12, optimizer variables = 5, metrics = 0, others = 197
# AWS-TF 2.1: all = 226, loss = 2, weights = 1, gradients = 2, biases = 12, optimizer variables = 5, metrics = 0, others = 204
assert len(tnames) >= 214
assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 1
assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 12
assert len(trial.tensor_names(collection=CollectionKeys.GRADIENTS)) >= 0
assert len(trial.tensor_names(collection=CollectionKeys.OPTIMIZER_VARIABLES)) >= 0
else:
assert len(tnames) == 2
assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 2
def test_save_layer_inputs_and_outputs(out_dir, tf_eager_mode):
# explicitly save INPUTS and OUTPUTS
include_collections = [CollectionKeys.INPUTS, CollectionKeys.OUTPUTS]
hook = smd.KerasHook(out_dir=out_dir,
include_collections=include_collections)
helper_keras_fit(
trial_dir=out_dir,
hook=hook,
eager=tf_eager_mode,
steps=["train", "eval", "predict", "train"],
)
trial = smd.create_trial(path=out_dir)
assert len(trial.tensor_names(collection=CollectionKeys.INPUTS)) == 4
assert len(trial.tensor_names(collection=CollectionKeys.OUTPUTS)) == 4
# Check that output of layer is equal to the input of the next
boolean_matrix = trial.tensor("flatten/outputs").value(0) == trial.tensor(
"dense/inputs").value(0)
assert boolean_matrix.all()
boolean_matrix = trial.tensor("dense/outputs").value(0) == trial.tensor(
"dropout/inputs").value(0)
assert boolean_matrix.all()
boolean_matrix = trial.tensor("dropout/outputs").value(0) == trial.tensor(
"dense_1/inputs").value(0)
assert boolean_matrix.all()
def test_gradtape_include_collections(out_dir):
"""
This test ensures that a training script written with GradientTape
handles the case where hook config contains all collections mentioned
through include collections
"""
include_collections = [
CollectionKeys.WEIGHTS,
CollectionKeys.BIASES,
CollectionKeys.GRADIENTS,
CollectionKeys.LOSSES,
CollectionKeys.OUTPUTS,
CollectionKeys.METRICS,
CollectionKeys.OPTIMIZER_VARIABLES,
]
save_config = SaveConfig(save_interval=3)
hook = smd.KerasHook(
out_dir,
save_config=save_config,
include_collections=include_collections,
reduction_config=ReductionConfig(norms=ALLOWED_NORMS,
reductions=ALLOWED_REDUCTIONS),
)
helper_keras_gradtape(out_dir, hook=hook)
trial = smd.create_trial(path=out_dir)
# can't save gradients in TF 2.x
assert len(trial.tensor_names()) == (16 if is_tf_2_2() else 15)
assert len(trial.tensor_names(collection=CollectionKeys.GRADIENTS)) == 4
assert len(
trial.tensor_names(collection=CollectionKeys.OPTIMIZER_VARIABLES)) == 5
assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 1
assert len(trial.tensor_names(collection=CollectionKeys.METRICS)) == 1
def test_linear_classifier(script_mode: bool):
""" Works as intended. """
smd.del_hook()
tf.reset_default_graph()
with SagemakerSimulator() as sim:
# Setup
train_input_fn, eval_input_fn = get_input_fns()
x_feature = tf.feature_column.numeric_column("x", shape=(28, 28))
estimator = tf.compat.v1.estimator.LinearClassifier(
feature_columns=[x_feature],
model_dir="/tmp/mnist_linear_classifier",
n_classes=10)
# Train
if script_mode:
hook = smd.EstimatorHook(out_dir=sim.out_dir)
estimator.train(input_fn=train_input_fn, steps=100, hooks=[hook])
else:
estimator.train(input_fn=train_input_fn, steps=100)
# 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."
def test_monitored_session(script_mode: bool):
""" Works as intended. """
smd.del_hook()
tf.reset_default_graph()
with SagemakerSimulator() as sim:
train_op, X, Y = get_train_op_and_placeholders()
init = tf.compat.v1.global_variables_initializer()
mnist = get_data()
if script_mode:
hook = smd.SessionHook(out_dir=sim.out_dir)
sess = tf.train.MonitoredSession(hooks=[hook])
else:
sess = tf.train.MonitoredSession()
with sess:
sess.run(init)
for step in range(1, 101):
batch_x, batch_y = mnist.train.next_batch(32)
sess.run(train_op, feed_dict={X: batch_x, Y: batch_y})
# 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."
def test_monitored_session(script_mode):
""" Works as intended. """
smd.del_hook()
tf.reset_default_graph()
json_file_contents = """
{
"S3OutputPath": "s3://sagemaker-test",
"LocalPath": "/opt/ml/output/tensors",
"HookParameters" : {
"save_interval": "100"
}
}
"""
with SagemakerSimulator(json_file_contents=json_file_contents) as sim:
train_op, X, Y = get_train_op_and_placeholders()
init = tf.global_variables_initializer()
mnist = get_data()
if script_mode:
hook = smd.SessionHook(out_dir=sim.out_dir)
sess = tf.train.MonitoredSession(hooks=[hook])
else:
sess = tf.train.MonitoredSession()
with sess:
sess.run(init)
for step in range(1, 101):
batch_x, batch_y = mnist.train.next_batch(32)
sess.run(train_op, feed_dict={X: batch_x, Y: batch_y})
# 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."
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."
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_fit_pure_eager(out_dir, tf_eager_mode):
"""
Test save all and save default collection in fit() pure eager mode
"""
hook = smd.KerasHook(out_dir=out_dir,
save_all=True,
save_config=SaveConfig(save_interval=3))
helper_keras_fit(trial_dir=out_dir,
hook=hook,
eager=tf_eager_mode,
run_eagerly=True)
trial = smd.create_trial(path=out_dir)
if is_tf_2_2():
assert len(trial.tensor_names()) == 27
else:
assert len(trial.tensor_names()) == (20 if is_tf_2_3() else 21)
assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 2
assert len(
trial.tensor_names(collection=CollectionKeys.OPTIMIZER_VARIABLES)) == 5
assert len(trial.tensor_names(
collection=CollectionKeys.INPUTS)) == (1 if is_tf_2_2() else 0)
assert len(trial.tensor_names(
collection=CollectionKeys.OUTPUTS)) == (2 if is_tf_2_2() else 0)
def test_keras_fit(out_dir, tf_eager_mode, saveall):
hook = smd.KerasHook(out_dir=out_dir, save_all=saveall)
helper_keras_fit(
trial_dir=out_dir,
hook=hook,
eager=tf_eager_mode,
steps=["train", "eval", "predict", "train"],
)
trial = smd.create_trial(path=out_dir)
# can't save gradients in TF 2.x eager mode
if saveall: # save losses, metrics, weights, biases
if tf_eager_mode:
assert len(trial.tensor_names()) == 7 if is_tf_2_2() else 8
else:
assert len(trial.tensor_names()) == 21
assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 2
else: # save the default losses and metrics
assert len(trial.tensor_names()) == 3 if is_tf_2_2() and tf_eager_mode else 4
assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 1
assert (
len(trial.tensor_names(collection=CollectionKeys.METRICS)) == 2
if is_tf_2_2() and tf_eager_mode
else 3
)
def test_estimator(script_mode: bool):
""" Works as intended. """
smd.del_hook()
tf.reset_default_graph()
with SagemakerSimulator() as sim:
# Setup
mnist_classifier = get_estimator()
train_input_fn, eval_input_fn = get_input_fns()
# Train and evaluate
train_steps, eval_steps = 80, 20
if script_mode:
hook = smd.EstimatorHook(out_dir=sim.out_dir)
hook.set_mode(mode=smd.modes.TRAIN)
mnist_classifier.train(input_fn=train_input_fn,
steps=train_steps,
hooks=[hook])
hook.set_mode(mode=smd.modes.EVAL)
mnist_classifier.evaluate(input_fn=eval_input_fn,
steps=eval_steps,
hooks=[hook])
else:
mnist_classifier.train(input_fn=train_input_fn, steps=train_steps)
mnist_classifier.evaluate(input_fn=eval_input_fn, steps=eval_steps)
# Check that hook created and tensors saved
trial = smd.create_trial(path=sim.out_dir)
print(trial)
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 trial.steps() == [0, train_steps], "Wrong step count for trial."
def test_model_inputs_and_outputs(out_dir, tf_eager_mode):
# explicitly save INPUTS and OUTPUTS
include_collections = [CollectionKeys.INPUTS, CollectionKeys.OUTPUTS]
hook = smd.KerasHook(out_dir=out_dir,
include_collections=include_collections)
helper_keras_fit(
trial_dir=out_dir,
hook=hook,
eager=tf_eager_mode,
steps=["train", "eval", "predict", "train"],
)
trial = smd.create_trial(path=out_dir)
assert len(trial.steps(mode=ModeKeys.TRAIN)) == 3
assert len(trial.tensor_names(collection=CollectionKeys.OUTPUTS)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.INPUTS)) == 1
for tname in trial.tensor_names(collection=CollectionKeys.OUTPUTS):
output = trial.tensor(tname)
assert tname in ["y", "y_pred"]
assert output.value(0) is not None
# Check the shape of output tensors
assert trial.tensor("y").value(0).shape[1] == 1 # label
assert trial.tensor("y_pred").value(
0).shape[1] == 10 # Output probability for each class
def test_include_collections(out_dir, tf_eager_mode):
include_collections = [
CollectionKeys.WEIGHTS,
CollectionKeys.BIASES,
CollectionKeys.GRADIENTS,
CollectionKeys.LOSSES,
CollectionKeys.OUTPUTS,
CollectionKeys.METRICS,
CollectionKeys.OPTIMIZER_VARIABLES,
]
save_config = SaveConfig(save_interval=3)
hook = smd.KerasHook(
out_dir,
save_config=save_config,
include_collections=include_collections,
reduction_config=ReductionConfig(norms=ALLOWED_NORMS, reductions=ALLOWED_REDUCTIONS),
)
helper_keras_fit(out_dir, hook=hook, steps=["train", "eval", "predict"], eager=tf_eager_mode)
trial = smd.create_trial(path=out_dir)
# can't save gradients in TF 2.x
if tf_eager_mode:
assert len(trial.tensor_names()) == 7 if is_tf_2_2() else 8
else:
assert len(trial.tensor_names()) == 18
assert len(trial.tensor_names(collection=CollectionKeys.GRADIENTS)) == 4
assert len(trial.tensor_names(collection=CollectionKeys.OPTIMIZER_VARIABLES)) == 5
assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 1
assert (
len(trial.tensor_names(collection=CollectionKeys.METRICS)) == 2
if is_tf_2_2() and tf_eager_mode
else 3
)
def test_gradtape_tf_function(out_dir):
def get_grads(images, labels):
# with tf.GradientTape() as tape:
return model(images, training=True)
@tf.function
def train_step(images, labels):
return tf.reduce_mean(get_grads(images, labels))
mnist = tf.keras.datasets.mnist
(x_train, y_train), _ = mnist.load_data()
dataset = tf.data.Dataset.from_tensor_slices(
(tf.cast(x_train[..., tf.newaxis] / 255,
tf.float32), tf.cast(y_train, tf.int64)))
dataset = dataset.shuffle(1000).batch(64)
model = create_model()
hook = create_hook(out_dir)
opt = tf.keras.optimizers.Adam()
hook.wrap_optimizer(opt)
n_epochs = 1
for epoch in range(n_epochs):
for data, labels in dataset:
dataset_labels = labels
labels = tf.one_hot(labels, depth=10)
with hook.wrap_tape(tf.GradientTape()) as tape:
logits = train_step(data, labels)
grads = tape.gradient(logits, model.variables)
opt.apply_gradients(zip(grads, model.variables))
hook.save_tensor("inputs", data, CollectionKeys.INPUTS)
hook.save_tensor("logits", logits, CollectionKeys.OUTPUTS)
hook.save_tensor("labels", labels, CollectionKeys.OUTPUTS)
model.save(out_dir, save_format="tf")
hook.close()
trial = smd.create_trial(out_dir)
assert trial.tensor_names(collection=CollectionKeys.LOSSES) == ["loss"]
assert trial.tensor_names(collection=CollectionKeys.WEIGHTS) == [
"weights/dense/kernel:0",
"weights/dense_1/kernel:0",
]
assert trial.tensor_names(collection=CollectionKeys.BIASES) == [
"weights/dense/bias:0",
"weights/dense_1/bias:0",
]
assert trial.tensor_names(
collection=CollectionKeys.OPTIMIZER_VARIABLES) == [
"Adam/beta_1:0",
"Adam/beta_2:0",
"Adam/decay:0",
"Adam/iter:0",
"Adam/learning_rate:0",
]
assert trial.tensor_names(collection=CollectionKeys.INPUTS) == ["inputs"]
assert trial.tensor_names(collection=CollectionKeys.OUTPUTS) == [
"labels", "logits"
]
def helper_test_keras_v2(script_mode: bool = False, eager_mode: bool = True):
""" Test the default ZCC behavior of saving losses and metrics in eager and non-eager modes."""
smd.del_hook()
tf.keras.backend.clear_session()
if not eager_mode:
tf.compat.v1.disable_eager_execution()
with SagemakerSimulator() as sim:
model = get_keras_model_v2()
(x_train, y_train), (x_test, y_test) = get_keras_data()
x_train, x_test = x_train / 255, x_test / 255
opt = tf.keras.optimizers.RMSprop()
if script_mode:
hook = smd.KerasHook(out_dir=sim.out_dir, export_tensorboard=True)
opt = hook.wrap_optimizer(opt)
model.compile(loss="sparse_categorical_crossentropy",
optimizer=opt,
metrics=["accuracy"])
history = model.fit(x_train,
y_train,
batch_size=64,
epochs=2,
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=64,
epochs=2,
validation_split=0.2)
test_scores = model.evaluate(x_test, y_test, verbose=2)
hook = smd.get_hook()
assert hook
hook.close()
# Check that hook created and tensors saved
trial = smd.create_trial(path=sim.out_dir)
assert len(trial.steps()) > 0, "Nothing saved at any step."
assert len(trial.tensor_names()) > 0, "Tensors were not saved."
# DEFAULT TENSORS SAVED
assert len(trial.tensor_names(
collection=CollectionKeys.LOSSES)) > 0, "No Losses Saved"
assert len(trial.tensor_names(
collection=CollectionKeys.METRICS)) > 0, "No Metrics Saved"
assert (len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 0
), "Weights were not expected to be saved by default"
assert (len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 0
), "Biases were not expected to be saved by default"
def helper_test_keras_v2_json_config(json_file_contents,
script_mode: bool = False,
eager_mode: bool = True):
""" Tests ZCC with custom hook configs """
smd.del_hook()
tf.keras.backend.clear_session()
if not eager_mode:
tf.compat.v1.disable_eager_execution()
with SagemakerSimulator(json_file_contents=json_file_contents) as sim:
model = get_keras_model_v2()
(x_train, y_train), (x_test, y_test) = get_keras_data()
x_train, x_test = x_train / 255, x_test / 255
opt = tf.keras.optimizers.RMSprop()
if script_mode:
hook = smd.KerasHook.create_from_json_file()
opt = hook.wrap_optimizer(opt)
model.compile(loss="sparse_categorical_crossentropy",
optimizer=opt,
metrics=["accuracy"])
history = model.fit(x_train,
y_train,
batch_size=64,
epochs=2,
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,
epochs=2,
batch_size=64,
validation_split=0.2)
test_scores = model.evaluate(x_test, y_test, verbose=2)
hook = smd.get_hook()
assert hook
hook.close()
# Check that hook created and tensors saved
trial = smd.create_trial(path=sim.out_dir)
assert len(trial.steps()) > 0, "Nothing saved at any step."
assert len(trial.tensor_names()) > 0, "Tensors were not saved."
if not eager_mode:
assert len(trial.tensor_names(collection="gradients")) > 0
assert len(trial.tensor_names(collection="weights")) > 0
assert len(trial.tensor_names(collection="losses")) > 0
def test_keras_to_estimator(script_mode):
""" Works as intended. """
import tensorflow.compat.v1.keras as keras
tf.reset_default_graph()
smd.del_hook()
keras.backend.clear_session()
with SagemakerSimulator() as sim:
model = keras.models.Sequential([
keras.layers.Dense(16, activation="relu", input_shape=(4, )),
keras.layers.Dropout(0.2),
keras.layers.Dense(1, activation="sigmoid"),
])
def input_fn():
split = tfds.Split.TRAIN
data_dir = TEST_DATASET_S3_PATH if use_s3_datasets() else None
dataset = tfds.load("iris",
data_dir=data_dir,
split=split,
as_supervised=True)
dataset = dataset.map(lambda features, labels:
({
"dense_input": features
}, labels))
dataset = dataset.batch(32).repeat()
return dataset
model.compile(loss="categorical_crossentropy", optimizer="adam")
model.summary()
keras_estimator = tf.keras.estimator.model_to_estimator(
keras_model=model, model_dir=sim.out_dir)
if script_mode:
hook = smd.EstimatorHook(sim.out_dir)
hook.set_mode(smd.modes.TRAIN)
keras_estimator.train(input_fn=input_fn, steps=25, hooks=[hook])
hook.set_mode(smd.modes.EVAL)
eval_result = keras_estimator.evaluate(input_fn=input_fn,
steps=10,
hooks=[hook])
else:
keras_estimator.train(input_fn=input_fn, steps=25)
keras_estimator.evaluate(input_fn=input_fn, steps=10)
tr = smd.create_trial(sim.out_dir)
assert len(tr.tensor_names()) == 1
assert tr.steps() == [0, 25]
assert len(tr.steps(smd.modes.TRAIN)) == 1
assert len(tr.steps(smd.modes.EVAL)) == 1
def test_keras_fit_false(out_dir, tf_eager_mode=False):
test_include_collections = [
CollectionKeys.LOSSES,
CollectionKeys.METRICS,
CollectionKeys.WEIGHTS,
CollectionKeys.BIASES,
CollectionKeys.GRADIENTS,
CollectionKeys.INPUTS,
CollectionKeys.OUTPUTS,
CollectionKeys.LAYERS,
CollectionKeys.OPTIMIZER_VARIABLES,
]
hook = smd.KerasHook(out_dir=out_dir,
include_collections=test_include_collections)
helper_keras_fit(
include_collections=test_include_collections,
trial_dir=out_dir,
hook=hook,
run_eagerly=tf_eager_mode,
steps=["train", "eval", "predict", "train"],
)
trial = smd.create_trial(path=out_dir)
# We first assert that none of the collections we requested for are empty
assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 1
assert len(trial.tensor_names(collection=CollectionKeys.METRICS)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.GRADIENTS)) == 4
assert len(trial.tensor_names(
collection=CollectionKeys.INPUTS)) == 1 # 1 Model Input
assert len(trial.tensor_names(
collection=CollectionKeys.OUTPUTS)) == 2 # 2 Model outputs
assert len(
trial.tensor_names(collection=CollectionKeys.OPTIMIZER_VARIABLES)) == 5
# We assert that all the tensors saved have a valid value
for tname in trial.tensor_names():
assert trial.tensor(tname).value(0) is not None
# We then analyse Layer Inputs and Layer Outputs
# Check that output of layer is equal to the input of the next
boolean_matrix = trial.tensor("flatten_1/outputs").value(
0) == trial.tensor("dense_2/inputs").value(0)
assert boolean_matrix.all()
boolean_matrix = trial.tensor("dense_2/outputs").value(0) == trial.tensor(
"dropout_1/inputs").value(0)
assert boolean_matrix.all()
boolean_matrix = trial.tensor("dropout_1/outputs").value(
0) == trial.tensor("dense_3/inputs").value(0)
assert boolean_matrix.all()
def test_keras_fit(out_dir, tf_eager_mode, saveall):
hook = smd.KerasHook(out_dir=out_dir, save_all=saveall)
ts = time.time()
hook.save_scalar("foobar", 1, sm_metric=True, timestamp=ts)
scalars_to_be_saved = dict()
scalars_to_be_saved["scalar/foobar"] = (ts, 0)
helper_keras_fit(
trial_dir=out_dir,
hook=hook,
run_eagerly=tf_eager_mode,
steps=["train", "eval", "predict", "train"],
)
trial = smd.create_trial(path=out_dir)
# can't save gradients in TF 2.x eager mode
if saveall: # save losses, metrics, weights, biases, scalar
if tf_eager_mode:
if is_tf_2_2():
assert len(trial.tensor_names()) == 28
else:
assert len(trial.tensor_names()) == (21 if is_tf_2_3() else 14)
assert len(trial.tensor_names(collection=CollectionKeys.INPUTS)) == (
1 if is_tf_2_2() else 0
)
assert len(trial.tensor_names(collection=CollectionKeys.OUTPUTS)) == (
2 if is_tf_2_2() else 0
)
else:
assert len(trial.tensor_names()) == 21
assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.OPTIMIZER_VARIABLES)) == 5
assert (
len(
trial.tensor_names(
collection=CollectionKeys.OPTIMIZER_VARIABLES, mode=ModeKeys.EVAL
)
)
== 0,
"No Optimizer Variables Should be Saved in EVAL Mode",
)
else: # save the default losses and metrics
assert len(trial.tensor_names()) == (
4 if (is_tf_2_2() or is_tf_2_3()) and tf_eager_mode else 5
)
assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 1
assert len(trial.tensor_names(collection=CollectionKeys.METRICS)) == (
2 if (is_tf_2_2() or is_tf_2_3()) and tf_eager_mode else 3
)
for tname in trial.tensor_names():
assert trial.tensor(tname).value(0) is not None
def test_include_only_custom_collection(out_dir, tf_eager_mode):
include_collections = ["custom_optimizer_variables"]
save_config = SaveConfig(save_interval=3)
hook = smd.KerasHook(
out_dir,
save_config=save_config,
include_collections=include_collections,
reduction_config=ReductionConfig(norms=ALLOWED_NORMS, reductions=ALLOWED_REDUCTIONS),
)
hook.get_collection("custom_optimizer_variables").include("Adam")
helper_keras_fit(out_dir, hook=hook, steps=["train", "eval", "predict"], eager=tf_eager_mode)
trial = smd.create_trial(path=out_dir)
assert len(trial.tensor_names()) == (8 if is_tf_2_2() and tf_eager_mode else 9)
assert len(trial.tensor_names(collection="custom_optimizer_variables")) == 5
def test_gradtape_persistent(out_dir, saveall):
"""
Test save all and save default collection
"""
hook = smd.KerasHook(out_dir=out_dir, save_all=saveall, save_config=SaveConfig(save_interval=3))
helper_keras_gradtape(trial_dir=out_dir, hook=hook, persistent=True)
trial = smd.create_trial(path=out_dir)
if saveall: # save losses, metrics, weights, biases
assert len(trial.tensor_names()) == 10
assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 2
else: # save the default losses and metrics
assert len(trial.tensor_names()) == 2
assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 1
assert len(trial.tensor_names(collection=CollectionKeys.METRICS)) == 1
def test_hook_from_json(out_dir, tf_eager_mode, monkeypatch):
monkeypatch.setenv(
CONFIG_FILE_PATH_ENV_STR,
"tests/tensorflow/hooks/test_json_configs/test_collection_defaults.json",
)
hook = smd.KerasHook.create_from_json_file()
helper_keras_fit(out_dir, hook=hook, steps=["train"], run_eagerly=tf_eager_mode)
trial = smd.create_trial(path=out_dir)
# can't save gradients in TF 2.x
assert len(trial.tensor_names()) == (5 if (is_tf_2_2() or is_tf_2_3()) and tf_eager_mode else 6)
assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 0
assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 1
assert len(trial.tensor_names(collection=CollectionKeys.METRICS)) == (
2 if (is_tf_2_2() or is_tf_2_3()) and tf_eager_mode else 3
)
def test_estimator(script_mode):
""" Works as intended. """
smd.del_hook()
tf.reset_default_graph()
with SagemakerSimulator() as sim:
train_steps, eval_steps = 80, 20
helper_train(
script_mode=script_mode, sim=sim, train_steps=train_steps, eval_steps=eval_steps
)
# Check that hook created and tensors saved
trial = smd.create_trial(path=sim.out_dir)
print(trial)
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 trial.steps() == [0, train_steps], "Wrong step count for trial."
def test_keras_gradtape(out_dir, saveall):
"""
Test save all and save default collection
"""
hook = smd.KerasHook(out_dir=out_dir, save_all=saveall, save_config=SaveConfig(save_interval=3))
helper_keras_gradtape(trial_dir=out_dir, hook=hook)
trial = smd.create_trial(path=out_dir)
if saveall: # save losses, metrics, weights, biases
assert len(trial.tensor_names()) == (25 if is_tf_2_2() else 15)
assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.OPTIMIZER_VARIABLES)) == 5
else: # save the default losses and metrics
assert len(trial.tensor_names()) == 2
assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 1
assert len(trial.tensor_names(collection=CollectionKeys.METRICS)) == 1
def test_save_gradients(out_dir, tf_eager_mode):
# explicitly save INPUTS and OUTPUTS
include_collections = [CollectionKeys.GRADIENTS]
hook = smd.KerasHook(out_dir=out_dir, include_collections=include_collections)
helper_keras_fit(
trial_dir=out_dir,
hook=hook,
eager=tf_eager_mode,
steps=["train", "eval", "predict", "train"],
)
trial = smd.create_trial(path=out_dir)
assert len(trial.tensor_names(collection=CollectionKeys.GRADIENTS)) == 4
for tname in trial.tensor_names(collection=CollectionKeys.GRADIENTS):
output = trial.tensor(tname)
assert output.value(0) is not None
def test_weights_collections(out_dir, tf_eager_mode):
hook = smd.KerasHook(
out_dir,
save_config=SaveConfig(save_interval=3),
include_collections=[CollectionKeys.WEIGHTS],
)
helper_keras_fit(out_dir, hook=hook, steps=["train"], run_eagerly=tf_eager_mode)
trial = smd.create_trial(path=out_dir)
# can't save gradients in TF 2.x
assert len(trial.tensor_names()) == (5 if (is_tf_2_2() or is_tf_2_3()) and tf_eager_mode else 6)
assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 0
assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 2
assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 1
assert len(trial.tensor_names(collection=CollectionKeys.METRICS)) == (
2 if (is_tf_2_2() or is_tf_2_3()) and tf_eager_mode else 3
)
