Minor improvements on examples.

PiperOrigin-RevId: 283798223

Author: csferng

neural_structured_learning/examples/adv_keras_cnn_mnist.py

@@ -37,6 +37,9 @@
 FLAGS = flags.FLAGS
 
 flags.DEFINE_integer('epochs', None, 'Number of epochs to train.')
+flags.DEFINE_integer('steps_per_epoch', None,
+                     'Number of steps in each training epoch.')
+flags.DEFINE_integer('eval_steps', None, 'Number of steps to evaluate.')
 flags.DEFINE_float('adv_step_size', None,
                    'Step size for generating adversarial examples.')
 
@@ -62,14 +65,21 @@ class HParams(object):
   batch_size = attr.ib(default=32)
   buffer_size = attr.ib(default=10000)
   epochs = attr.ib(default=5)
+  steps_per_epoch = attr.ib(default=None)
+  eval_steps = attr.ib(default=None)
 
 
 def get_hparams():
+  """Returns the hyperparameters with defaults overwritten by flags."""
   hparams = HParams()
   if FLAGS.epochs:
     hparams.epochs = FLAGS.epochs
   if FLAGS.adv_step_size:
     hparams.adv_step_size = FLAGS.adv_step_size
+  if FLAGS.steps_per_epoch:
+    hparams.steps_per_epoch = FLAGS.steps_per_epoch
+  if FLAGS.eval_steps:
+    hparams.eval_steps = FLAGS.eval_steps
   return hparams
 
 
@@ -136,8 +146,11 @@ def train_and_evaluate(model, hparams, train_dataset, test_dataset):
       optimizer='adam',
       loss='sparse_categorical_crossentropy',
       metrics=['accuracy'])
-  model.fit(train_dataset, epochs=hparams.epochs)
-  eval_result = model.evaluate(test_dataset)
+  model.fit(
+      train_dataset,
+      epochs=hparams.epochs,
+      steps_per_epoch=hparams.steps_per_epoch)
+  eval_result = model.evaluate(test_dataset, steps=hparams.eval_steps)
   return list(zip(model.metrics_names, eval_result))
 
 
@@ -169,6 +182,8 @@ def evaluate_robustness(model_to_attack, dataset, models, hparams):
       for name in models.keys()
   }
 
+  if hparams.eval_steps:
+    dataset = dataset.take(hparams.eval_steps)
   # When running on accelerators, looping over the dataset inside a tf.function
   # may be much faster.
   for batch in dataset:

neural_structured_learning/examples/graph_keras_mlp_cora.py

@@ -43,6 +43,9 @@
 FLAGS = flags.FLAGS
 FLAGS.showprefixforinfo = False
 
+flags.DEFINE_integer('train_epochs', None, 'Number of epochs to train.')
+flags.DEFINE_integer('eval_steps', None, 'Number of steps to evaluate.')
+
 NBR_FEATURE_PREFIX = 'NL_nbr_'
 NBR_WEIGHT_SUFFIX = '_weight'
 
@@ -69,7 +72,12 @@ class HParams(object):
 
 def get_hyper_parameters():
   """Returns the hyper-parameters used for training."""
-  return HParams()
+  hparams = HParams()
+  if FLAGS.train_epochs:
+    hparams.train_epochs = FLAGS.train_epochs
+  if FLAGS.eval_steps:
+    hparams.eval_steps = FLAGS.eval_steps
+  return hparams
 
 
 def load_dataset(filename):

neural_structured_learning/examples/preprocess/cora/prep_data.sh

@@ -38,7 +38,7 @@ tar -C ${DATA_DIR} -xvzf ${DATA_DIR}/cora.tgz
 
 # Pre-process cora dataset. The file 'preprocess_cora_dataset.py' is assumed to
 # be located in the current directory.
-python preprocess_cora_dataset.py \
+python $(dirname "$0")/preprocess_cora_dataset.py \
 --input_cora_content=${DATA_DIR}/cora/cora.content \
 --input_cora_graph=${DATA_DIR}/cora/cora.cites \
 --max_nbrs=5 \