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 )