Skip to content

Commit 0afd72f

Browse files
NihalHarishNihalHarish
authored
Save Model Inputs, Model Outputs, Gradients, Custom Tensors, Layer Inputs, Layer Outputs (aws#282)
1 parent 58ab100 commit 0afd72f

26 files changed

+940
-42
lines changed

‎examples/tensorflow2/scripts/tf2_save_metrics.py

Lines changed: 209 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,209 @@
1+
"""
2+
This file is temporary, for testing with 2.X.
3+
We'll need to integrate a more robust testing pipeline and make this part of pytest
4+
before pushing to master.
5+
6+
This was tested with TensorFlow 2.1, by running
7+
`python tests/tensorflow2/test_keras.py` from the main directory.
8+
"""
9+
# Standard Library
10+
import shutil
11+
12+
# Third Party
13+
import pytest
14+
import tensorflow.compat.v2 as tf
15+
16+
# First Party
17+
import smdebug.tensorflow as smd
18+
from smdebug.core.collection import CollectionKeys
19+
from smdebug.tensorflow import SaveConfig
20+
21+
22+
@pytest.fixture(scope="function")
23+
def out_dir():
24+
""" Use this method to construct an out_dir.
25+
26+
Then it will be automatically cleaned up for you, passed into the test method, and we'll have
27+
fewer folders lying around.
28+
"""
29+
out_dir = "/tmp/test"
30+
shutil.rmtree(out_dir, ignore_errors=True)
31+
return out_dir
32+
33+
34+
def helper_keras_fit(
35+
trial_dir,
36+
save_all=False,
37+
include_collections=None,
38+
reduction_config=None,
39+
save_config=None,
40+
hook=None,
41+
steps=None,
42+
add_callbacks=None,
43+
run_eagerly=False,
44+
):
45+
46+
mnist = tf.keras.datasets.mnist
47+
(x_train, y_train), (x_test, y_test) = mnist.load_data()
48+
x_train, x_test = x_train / 255, x_test / 255
49+
50+
model = tf.keras.models.Sequential(
51+
[
52+
tf.keras.layers.Flatten(input_shape=(28, 28)),
53+
tf.keras.layers.Dense(128, activation="relu"),
54+
tf.keras.layers.Dropout(0.2),
55+
tf.keras.layers.Dense(10, activation="softmax"),
56+
]
57+
)
58+
59+
if hook is None:
60+
if save_config is None:
61+
save_config = SaveConfig(save_interval=3)
62+
63+
hook = smd.KerasHook(
64+
trial_dir,
65+
save_config=save_config,
66+
save_all=save_all,
67+
include_collections=include_collections,
68+
reduction_config=reduction_config,
69+
)
70+
71+
if not save_all and include_collections is not None:
72+
for cname in hook.include_collections:
73+
if cname not in include_collections:
74+
hook.get_collection(cname).save_config = SaveConfig(end_step=0)
75+
76+
opt = tf.keras.optimizers.Adam()
77+
78+
opt = hook.wrap_optimizer(opt)
79+
model.compile(
80+
optimizer=opt,
81+
loss="sparse_categorical_crossentropy",
82+
metrics=["accuracy"],
83+
run_eagerly=run_eagerly,
84+
)
85+
hooks = []
86+
if add_callbacks:
87+
if "tensorboard" in add_callbacks:
88+
hooks.append(
89+
tf.keras.callbacks.TensorBoard(
90+
log_dir="/tmp/logs", histogram_freq=1, write_grads=True, write_images=True
91+
)
92+
)
93+
hooks.append(hook)
94+
95+
if steps is None:
96+
steps = ["train"]
97+
for step in steps:
98+
if step == "train":
99+
model.fit(x_train, y_train, epochs=1, steps_per_epoch=10, callbacks=hooks, verbose=0)
100+
elif step == "eval":
101+
model.evaluate(x_test, y_test, steps=10, callbacks=hooks, verbose=0)
102+
elif step == "predict":
103+
model.predict(x_test[:100], callbacks=hooks, verbose=0)
104+
105+
hook.close()
106+
107+
108+
def test_keras_fit_eager(out_dir, tf_eager_mode=True):
109+
test_include_collections = [
110+
CollectionKeys.LOSSES,
111+
CollectionKeys.METRICS,
112+
CollectionKeys.WEIGHTS,
113+
CollectionKeys.BIASES,
114+
CollectionKeys.GRADIENTS,
115+
CollectionKeys.INPUTS,
116+
CollectionKeys.OUTPUTS,
117+
CollectionKeys.LAYERS,
118+
CollectionKeys.OPTIMIZER_VARIABLES,
119+
]
120+
hook = smd.KerasHook(out_dir=out_dir, include_collections=test_include_collections)
121+
helper_keras_fit(
122+
include_collections=test_include_collections,
123+
trial_dir=out_dir,
124+
hook=hook,
125+
run_eagerly=tf_eager_mode,
126+
steps=["train", "eval", "predict", "train"],
127+
)
128+
trial = smd.create_trial(path=out_dir)
129+
130+
# We first assert that none of the collections we requested for are empty
131+
assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 1
132+
assert len(trial.tensor_names(collection=CollectionKeys.METRICS)) == 2
133+
assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 2
134+
assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 2
135+
assert len(trial.tensor_names(collection=CollectionKeys.GRADIENTS)) == 4
136+
assert len(trial.tensor_names(collection=CollectionKeys.INPUTS)) == 1 # 1 Model Input
137+
assert len(trial.tensor_names(collection=CollectionKeys.OUTPUTS)) == 2 # 2 Model outputs
138+
assert len(trial.tensor_names(collection=CollectionKeys.OPTIMIZER_VARIABLES)) == 5
139+
140+
# We assert that all the tensors saved have a valid value
141+
for tname in trial.tensor_names():
142+
assert trial.tensor(tname).value(0) is not None
143+
144+
# We then analyse Layer Inputs and Layer Outputs
145+
# Check that output of layer is equal to the input of the next
146+
boolean_matrix = trial.tensor("flatten/outputs").value(0) == trial.tensor("dense/inputs").value(
147+
0
148+
)
149+
assert boolean_matrix.all()
150+
boolean_matrix = trial.tensor("dense/outputs").value(0) == trial.tensor("dropout/inputs").value(
151+
0
152+
)
153+
assert boolean_matrix.all()
154+
boolean_matrix = trial.tensor("dropout/outputs").value(0) == trial.tensor(
155+
"dense_1/inputs"
156+
).value(0)
157+
assert boolean_matrix.all()
158+
159+
160+
def test_keras_fit_false(out_dir, tf_eager_mode=False):
161+
test_include_collections = [
162+
CollectionKeys.LOSSES,
163+
CollectionKeys.METRICS,
164+
CollectionKeys.WEIGHTS,
165+
CollectionKeys.BIASES,
166+
CollectionKeys.GRADIENTS,
167+
CollectionKeys.INPUTS,
168+
CollectionKeys.OUTPUTS,
169+
CollectionKeys.LAYERS,
170+
CollectionKeys.OPTIMIZER_VARIABLES,
171+
]
172+
hook = smd.KerasHook(out_dir=out_dir, include_collections=test_include_collections)
173+
helper_keras_fit(
174+
include_collections=test_include_collections,
175+
trial_dir=out_dir,
176+
hook=hook,
177+
run_eagerly=tf_eager_mode,
178+
steps=["train", "eval", "predict", "train"],
179+
)
180+
trial = smd.create_trial(path=out_dir)
181+
182+
# We first assert that none of the collections we requested for are empty
183+
assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 1
184+
assert len(trial.tensor_names(collection=CollectionKeys.METRICS)) == 2
185+
assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 2
186+
assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 2
187+
assert len(trial.tensor_names(collection=CollectionKeys.GRADIENTS)) == 4
188+
assert len(trial.tensor_names(collection=CollectionKeys.INPUTS)) == 1 # 1 Model Input
189+
assert len(trial.tensor_names(collection=CollectionKeys.OUTPUTS)) == 2 # 2 Model outputs
190+
assert len(trial.tensor_names(collection=CollectionKeys.OPTIMIZER_VARIABLES)) == 5
191+
192+
# We assert that all the tensors saved have a valid value
193+
for tname in trial.tensor_names():
194+
assert trial.tensor(tname).value(0) is not None
195+
196+
# We then analyse Layer Inputs and Layer Outputs
197+
# Check that output of layer is equal to the input of the next
198+
boolean_matrix = trial.tensor("flatten_1/outputs").value(0) == trial.tensor(
199+
"dense_2/inputs"
200+
).value(0)
201+
assert boolean_matrix.all()
202+
boolean_matrix = trial.tensor("dense_2/outputs").value(0) == trial.tensor(
203+
"dropout_1/inputs"
204+
).value(0)
205+
assert boolean_matrix.all()
206+
boolean_matrix = trial.tensor("dropout_1/outputs").value(0) == trial.tensor(
207+
"dense_3/inputs"
208+
).value(0)
209+
assert boolean_matrix.all()

‎examples/tensorflow2/scripts/tf_keras_gradienttape.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,9 @@ def train(batch_size, n_epochs, model, hook):
6565
optimizer.apply_gradients(zip(grads, model.trainable_variables))
6666
acc = train_acc_metric(dataset_labels, logits)
6767
# save metrics value
68-
hook.record_tensor_value(tensor_name="accuracy", tensor_value=acc)
68+
hook.save_tensor(
69+
tensor_name="accuracy", tensor_value=acc, collections_to_write="metrics"
70+
)
6971
values = [("Accuracy", train_acc_metric.result())]
7072
progBar.update(idx * batch_size, values=values)
7173

‎examples/tensorflow2/scripts/tf_save_metrics_gradient_tape.py

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
"""
2+
This file is temporary, for testing with 2.X.
3+
We'll need to integrate a more robust testing pipeline and make this part of pytest
4+
before pushing to master.
5+
"""
6+
# Standard Library
7+
import shutil
8+
9+
# Third Party
10+
import pytest
11+
import tensorflow.compat.v2 as tf
12+
13+
# First Party
14+
import smdebug.tensorflow as smd
15+
from smdebug.core.collection import CollectionKeys
16+
from smdebug.tensorflow import SaveConfig
17+
18+
19+
@pytest.fixture(scope="function")
20+
def out_dir():
21+
""" Use this method to construct an out_dir.
22+
23+
Then it will be automatically cleaned up for you, passed into the test method, and we'll have
24+
fewer folders lying around.
25+
"""
26+
out_dir = "/tmp/test"
27+
shutil.rmtree(out_dir, ignore_errors=True)
28+
return out_dir
29+
30+
31+
def helper_keras_gradtape(
32+
trial_dir,
33+
save_all=False,
34+
include_collections=None,
35+
reduction_config=None,
36+
save_config=None,
37+
hook=None,
38+
batch_size=64,
39+
persistent=False,
40+
):
41+
mnist = tf.keras.datasets.mnist
42+
(x_train, y_train), _ = mnist.load_data()
43+
dataset = tf.data.Dataset.from_tensor_slices(
44+
(tf.cast(x_train[..., tf.newaxis] / 255, tf.float32), tf.cast(y_train, tf.int64))
45+
)
46+
dataset = dataset.shuffle(1000).batch(batch_size)
47+
48+
model = tf.keras.models.Sequential(
49+
[
50+
# WA for TF issue https://github.com/tensorflow/tensorflow/issues/36279
51+
tf.keras.layers.Flatten(input_shape=(28, 28, 1)),
52+
tf.keras.layers.Dense(128, activation="relu"),
53+
tf.keras.layers.Dropout(0.2),
54+
tf.keras.layers.Dense(10, activation="softmax"),
55+
]
56+
)
57+
58+
if hook is None:
59+
if save_config is None:
60+
save_config = SaveConfig(save_interval=3)
61+
62+
hook = smd.KerasHook(
63+
trial_dir,
64+
save_config=save_config,
65+
save_all=save_all,
66+
include_collections=include_collections,
67+
reduction_config=reduction_config,
68+
)
69+
70+
if not save_all and include_collections is not None:
71+
for cname in hook.include_collections:
72+
if cname not in include_collections:
73+
hook.get_collection(cname).save_config = SaveConfig(end_step=0)
74+
75+
opt = tf.keras.optimizers.Adam()
76+
hook.wrap_optimizer(opt)
77+
hook.register_model(model) # Can be skipped in ZCC
78+
79+
cce = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
80+
train_acc_metric = tf.keras.metrics.SparseCategoricalAccuracy()
81+
82+
n_epochs = 1
83+
for epoch in range(n_epochs):
84+
for data, labels in dataset:
85+
dataset_labels = labels
86+
labels = tf.one_hot(labels, depth=10)
87+
with hook.wrap_tape(tf.GradientTape(persistent=persistent)) as tape:
88+
logits = model(data, training=True)
89+
loss_value = cce(labels, logits)
90+
hook.save_tensor("y_labels", labels, "outputs")
91+
grads = tape.gradient(loss_value, model.variables)
92+
93+
# By default, the resources held by a GradientTape are released as
94+
# soon as GradientTape.gradient() method is called. To compute
95+
# multiple gradients over the same computation, create a persistent
96+
# gradient tape. This allows multiple calls to the gradient() method
97+
# as resources are released when the tape object is garbage collected.
98+
if persistent:
99+
_ = tape.gradient(loss_value, model.variables)
100+
opt.apply_gradients(zip(grads, model.variables))
101+
acc = train_acc_metric(dataset_labels, logits)
102+
hook.save_tensor(
103+
tensor_name="accuracy",
104+
tensor_value=acc,
105+
collections_to_write=CollectionKeys.METRICS,
106+
)
107+
train_acc_metric.reset_states()
108+
109+
hook.close()
110+
111+
112+
def test_keras_gradtape(out_dir):
113+
"""
114+
Test save all and save default collection
115+
"""
116+
include_collections = [
117+
CollectionKeys.WEIGHTS,
118+
CollectionKeys.BIASES,
119+
CollectionKeys.GRADIENTS,
120+
CollectionKeys.LAYERS,
121+
CollectionKeys.LOSSES,
122+
CollectionKeys.INPUTS,
123+
CollectionKeys.OUTPUTS,
124+
CollectionKeys.METRICS,
125+
CollectionKeys.OPTIMIZER_VARIABLES,
126+
]
127+
hook = smd.KerasHook(
128+
out_dir=out_dir,
129+
save_config=SaveConfig(save_interval=1),
130+
include_collections=include_collections,
131+
)
132+
helper_keras_gradtape(trial_dir=out_dir, hook=hook)
133+
134+
trial = smd.create_trial(path=out_dir)
135+
assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 2
136+
assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 2
137+
assert len(trial.tensor_names(collection=CollectionKeys.OPTIMIZER_VARIABLES)) == 5
138+
assert len(trial.tensor_names(collection=CollectionKeys.LAYERS)) == 8
139+
assert len(trial.tensor_names(collection=CollectionKeys.OUTPUTS)) == 2
140+
assert len(trial.tensor_names(collection=CollectionKeys.INPUTS)) == 1
141+
assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 1
142+
assert len(trial.tensor_names(collection=CollectionKeys.METRICS)) == 1
143+
144+
# We assert that all the tensors saved have a valid value
145+
for tname in trial.tensor_names():
146+
assert trial.tensor(tname).value(5) is not None
147+
148+
# We then analyse Layer Inputs and Layer Outputs
149+
# Check that output of a layer is equal to the input of the next
150+
boolean_matrix = trial.tensor("flatten/outputs").value(0) == trial.tensor("dense/inputs").value(
151+
0
152+
)
153+
assert boolean_matrix.all()
154+
boolean_matrix = trial.tensor("dense/outputs").value(0) == trial.tensor("dropout/inputs").value(
155+
0
156+
)
157+
assert boolean_matrix.all()
158+
boolean_matrix = trial.tensor("dropout/outputs").value(0) == trial.tensor(
159+
"dense_1/inputs"
160+
).value(0)
161+
assert boolean_matrix.all()

‎smdebug/core/collection.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ class CollectionKeys:
2626
GRADIENTS = "gradients"
2727
LOSSES = "losses"
2828
BIASES = "biases"
29+
LAYERS = "layers"
2930

3031
# Use this collection to log scalars other than losses/metrics to SageMaker.
3132
# Mainly for Tensorflow. For all other frameworks, call save_scalar() API
@@ -75,6 +76,7 @@ class CollectionKeys:
7576
CollectionKeys.METRICS,
7677
CollectionKeys.INPUTS,
7778
CollectionKeys.OUTPUTS,
79+
CollectionKeys.LAYERS,
7880
CollectionKeys.SM_METRICS,
7981
CollectionKeys.OPTIMIZER_VARIABLES,
8082
}

0 commit comments

Comments
 (0)