Add training history for visualization in image classification.

ziyeqinghan · copybara-github · commit 57a32876ed30 · 2020-04-21T23:35:49.000-07:00
PiperOrigin-RevId: 307756398
diff --git a/tensorflow_examples/lite/model_maker/core/task/image_classifier.py b/tensorflow_examples/lite/model_maker/core/task/image_classifier.py
@@ -158,6 +158,7 @@ def __init__(self,
         self.model_spec.mean_rgb,
         self.model_spec.stddev_rgb,
         use_augmentation=use_augmentation)
+    self.history = None  # Training history that returns from `keras_model.fit`.
 
   def _create_model(self, hparams=None):
     """Creates the classifier model for retraining."""
@@ -199,8 +200,9 @@ def train(self, train_data, validation_data=None, hparams=None):
     lib = hub_lib
     if isinstance(hparams, train_image_classifier_lib.HParams):
       lib = train_image_classifier_lib
-    return lib.train_model(self.model, hparams, train_data_and_size,
-                           validation_data_and_size)
+    self.history = lib.train_model(self.model, hparams, train_data_and_size,
+                                   validation_data_and_size)
+    return self.history
 
   def preprocess(self, image, label, is_training=False):
     return self.preprocessor(image, label, is_training)
