Model Maker: Set DEFAULT_EXPORT_FORMAT and ALLOWED_EXPORT_FORMAT for …

…base class. PiperOrigin-RevId: 339408350
tensorflow · Oct 28, 2020 · ae10e14 · ae10e14
1 parent e17ed53
commit ae10e14
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/tensorflow_examples/lite/model_maker/core/task/classification_model.py b/tensorflow_examples/lite/model_maker/core/task/classification_model.py
@@ -21,13 +21,18 @@
 import numpy as np
 import tensorflow.compat.v2 as tf
 from tensorflow_examples.lite.model_maker.core.data_util import data_util
+from tensorflow_examples.lite.model_maker.core.export_format import ExportFormat
 from tensorflow_examples.lite.model_maker.core.task import custom_model
 from tensorflow_examples.lite.model_maker.core.task import model_util
 
 
 class ClassificationModel(custom_model.CustomModel):
   """"The abstract base class that represents a Tensorflow classification model."""
 
+  DEFAULT_EXPORT_FORMAT = (ExportFormat.TFLITE, ExportFormat.LABEL)
+  ALLOWED_EXPORT_FORMAT = (ExportFormat.TFLITE, ExportFormat.LABEL,
+                           ExportFormat.SAVED_MODEL)
+
   def __init__(self, model_spec, index_to_label, num_classes, shuffle,
                train_whole_model):
     """Initialize a instance with data, deploy mode and other related parameters.

diff --git a/tensorflow_examples/lite/model_maker/core/task/custom_model.py b/tensorflow_examples/lite/model_maker/core/task/custom_model.py
@@ -29,8 +29,8 @@
 class CustomModel(abc.ABC):
   """"The abstract base class that represents a Tensorflow classification model."""
 
-  DEFAULT_EXPORT_FORMAT = ()
-  ALLOWED_EXPORT_FORMAT = ()
+  DEFAULT_EXPORT_FORMAT = (ExportFormat.TFLITE)
+  ALLOWED_EXPORT_FORMAT = (ExportFormat.TFLITE, ExportFormat.SAVED_MODEL)
 
   def __init__(self, model_spec, shuffle):
     """Initialize a instance with data, deploy mode and other related parameters.