Add basic error checking for log_dir for PruningSummaries.

PiperOrigin-RevId: 316778704

Author: alanchiao

tensorflow_model_optimization/python/core/sparsity/keras/BUILD

@@ -105,6 +105,7 @@ py_library(
     deps = [
         ":pruning_wrapper",
         # numpy dep1,
+        # six dep1,
         # tensorflow dep1,
         "//tensorflow_model_optimization/python/core/keras:compat",
     ],

tensorflow_model_optimization/python/core/sparsity/keras/pruning_callbacks.py

@@ -20,6 +20,7 @@
 
 # import g3
 import numpy as np
+import six
 import tensorflow as tf
 
 from tensorflow_model_optimization.python.core.keras import compat
@@ -95,6 +96,11 @@ class PruningSummaries(callbacks.TensorBoard):
   """
 
   def __init__(self, log_dir, update_freq='epoch', **kwargs):
+    if not isinstance(log_dir, six.string_types) or not log_dir:
+      raise ValueError(
+          '`log_dir` must be a non-empty string. You passed `log_dir`='
+          '{input}.'.format(input=log_dir))
+
     super(PruningSummaries, self).__init__(
         log_dir=log_dir, update_freq=update_freq, **kwargs)
 

tensorflow_model_optimization/python/core/sparsity/keras/pruning_callbacks_test.py

@@ -137,6 +137,17 @@ def testPruneTrainingLoopRaisesError_PruningStepCallbackMissing_CustomTrainingLo
       with tf.GradientTape():
         pruned_model(inp, training=True)
 
+  @keras_parameterized.run_all_keras_modes
+  def testPruningSummariesRaisesError_LogDirNotNonEmptyString(self):
+    with self.assertRaises(ValueError):
+      pruning_callbacks.PruningSummaries(log_dir='')
+
+    with self.assertRaises(ValueError):
+      pruning_callbacks.PruningSummaries(log_dir=None)
+
+    with self.assertRaises(ValueError):
+      pruning_callbacks.PruningSummaries(log_dir=object())
+
 
 if __name__ == '__main__':
   tf.test.main()