fix: correct Debugger tests (aws#305)

Author: knakad

tests/data/mxnet_mnist/mnist_gluon.py

@@ -0,0 +1,136 @@
+# Standard Library
+import argparse
+import random
+
+# Third Party
+import mxnet as mx
+import numpy as np
+from mxnet import autograd, gluon
+from mxnet.gluon import nn
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Train a mxnet gluon model for FashonMNIST dataset"
+    )
+    parser.add_argument("--batch-size", type=int, default=256, help="Batch size")
+    parser.add_argument("--epochs", type=int, default=1, help="Number of Epochs")
+    parser.add_argument("--learning_rate", type=float, default=0.1)
+    parser.add_argument(
+        "--context", type=str, default="cpu", help="Context can be either cpu or gpu"
+    )
+    parser.add_argument(
+        "--validate", type=bool, default=True, help="Run validation if running with smdebug"
+    )
+
+    opt = parser.parse_args()
+    return opt
+
+
+def test(ctx, net, val_data):
+    metric = mx.metric.Accuracy()
+    for i, (data, label) in enumerate(val_data):
+        data = data.as_in_context(ctx)
+        label = label.as_in_context(ctx)
+        output = net(data)
+        metric.update([label], [output])
+
+    return metric.get()
+
+
+def train_model(net, epochs, ctx, learning_rate, momentum, train_data, val_data):
+    # Collect all parameters from net and its children, then initialize them.
+    net.initialize(mx.init.Xavier(magnitude=2.24), ctx=ctx)
+    # Trainer is for updating parameters with gradient.
+    trainer = gluon.Trainer(
+        net.collect_params(), "sgd", {"learning_rate": learning_rate, "momentum": momentum}
+    )
+    metric = mx.metric.Accuracy()
+    loss = gluon.loss.SoftmaxCrossEntropyLoss()
+
+    for epoch in range(epochs):
+        # reset data iterator and metric at begining of epoch.
+        metric.reset()
+        for i, (data, label) in enumerate(train_data):
+            # Copy data to ctx if necessary
+            data = data.as_in_context(ctx)
+            label = label.as_in_context(ctx)
+            # Start recording computation graph with record() section.
+            # Recorded graphs can then be differentiated with backward.
+            with autograd.record():
+                output = net(data)
+                L = loss(output, label)
+                L.backward()
+            # take a gradient step with batch_size equal to data.shape[0]
+            trainer.step(data.shape[0])
+            # update metric at last.
+            metric.update([label], [output])
+
+            if i % 100 == 0 and i > 0:
+                name, acc = metric.get()
+                print("[Epoch %d Batch %d] Training: %s=%f" % (epoch, i, name, acc))
+
+        name, acc = metric.get()
+        print("[Epoch %d] Training: %s=%f" % (epoch, name, acc))
+        name, val_acc = test(ctx, net, val_data)
+        print("[Epoch %d] Validation: %s=%f" % (epoch, name, val_acc))
+
+
+def transformer(data, label):
+    data = data.reshape((-1,)).astype(np.float32) / 255
+    return data, label
+
+
+def prepare_data(batch_size):
+    train_data = gluon.data.DataLoader(
+        gluon.data.vision.MNIST("/tmp", train=True, transform=transformer),
+        batch_size=batch_size,
+        shuffle=True,
+        last_batch="discard",
+    )
+
+    val_data = gluon.data.DataLoader(
+        gluon.data.vision.MNIST("/tmp", train=False, transform=transformer),
+        batch_size=batch_size,
+        shuffle=False,
+    )
+    return train_data, val_data
+
+
+# Create a model using gluon API. The hook is currently
+# supports MXNet gluon models only.
+def create_gluon_model():
+    net = nn.Sequential()
+    with net.name_scope():
+        net.add(nn.Dense(128, activation="relu"))
+        net.add(nn.Dense(64, activation="relu"))
+        net.add(nn.Dense(10))
+    return net
+
+
+def main():
+    opt = parse_args()
+    mx.random.seed(128)
+    random.seed(12)
+    np.random.seed(2)
+
+    context = mx.cpu() if opt.context.lower() == "cpu" else mx.gpu()
+    # Create a Gluon Model.
+    net = create_gluon_model()
+
+    # Start the training.
+    train_data, val_data = prepare_data(opt.batch_size)
+
+    train_model(
+        net=net,
+        epochs=opt.epochs,
+        ctx=context,
+        learning_rate=opt.learning_rate,
+        momentum=0.9,
+        train_data=train_data,
+        val_data=val_data,
+    )
+
+
+if __name__ == "__main__":
+    main()

tests/data/mxnet_mnist/my_custom_rule.py

@@ -0,0 +1,14 @@
+from smdebug.rules.rule import Rule
+
+
+class CustomGradientRule(Rule):
+    def __init__(self, base_trial, threshold=10.0):
+        super().__init__(base_trial)
+        self.threshold = float(threshold)
+
+    def set_required_tensors(self, step):
+        for tname in self.base_trial.tensor_names(collection="gradients"):
+            self.req_tensors.add(tname, steps=[step])
+
+    def invoke_at_step(self, step):
+        return False