Update wide_n_deep_tutorial.py (tensorflow#11722)

alanyee · Vijay Vasudevan · commit 41881b93b1a2 · 2017-07-24T17:11:25.000-07:00
-mkdtemp is low-level, and creates a temporary file somewhere out of sight that requires users to delete the file themselves. For the use case of a tutorial, it is best to use a method that cleanups the file when finished
-TemporaryFile does such as described above, but not available in Python 2
diff --git a/tensorflow/examples/learn/wide_n_deep_tutorial.py b/tensorflow/examples/learn/wide_n_deep_tutorial.py
@@ -18,6 +18,7 @@
 from __future__ import print_function
 
 import argparse
+import shutil
 import sys
 import tempfile
 
@@ -175,6 +176,7 @@ def input_fn(data_file, num_epochs, shuffle):
 def train_and_eval(model_dir, model_type, train_steps, train_data, test_data):
   """Train and evaluate the model."""
   train_file_name, test_file_name = maybe_download(train_data, test_data)
+  # Specify file path below if want to find the output easily
   model_dir = tempfile.mkdtemp() if not model_dir else model_dir
 
   m = build_estimator(model_dir, model_type)
@@ -189,6 +191,8 @@ def train_and_eval(model_dir, model_type, train_steps, train_data, test_data):
   print("model directory = %s" % model_dir)
   for key in sorted(results):
     print("%s: %s" % (key, results[key]))
+  # Manual cleanup
+  shutil.rmtree(model_dir)
 
 
 FLAGS = None
