First commit XiUOS

2021-04-28 17:49:18 +08:00
commit 6001051eb7
1331 changed files with 433955 additions and 0 deletions
@@ -0,0 +1,120 @@
+# ==========================================================================================
+# Copyright (c) 2020 AIIT XUOS Lab
+# XiOS is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
+#        http://license.coscl.org.cn/MulanPSL2
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# ==========================================================================================
+#!/usr/bin/env python3
+
+import os
+import tensorflow as tf
+
+print("TensorFlow version %s" % (tf.__version__))
+
+MODEL_NAME_H5 = 'mnist.h5'
+MODEL_NAME_TFLITE = 'mnist.tflite'
+DEFAULT_QUAN_MODEL_NAME_TFLITE = 'mnist-default-quan.tflite'
+FULL_QUAN_MODEL_NAME_TFLITE = 'mnist-full-quan.tflite'
+
+def build_model(model_name):
+    print('\n>>> load mnist dataset')
+    mnist = tf.keras.datasets.mnist
+    (train_images, train_labels),(test_images, test_labels) = mnist.load_data()
+    print("train images shape: ", train_images.shape)
+    print("train labels shape: ", train_labels.shape)
+    print("test images shape: ", test_images.shape)
+    print("test labels shape: ", test_labels.shape)
+
+    # transform label to categorical, like: 2 -> [0, 0, 1, 0, 0, 0, 0, 0, 0, 0]
+    print('\n>>> transform label to categorical')
+    train_labels = tf.keras.utils.to_categorical(train_labels)
+    test_labels = tf.keras.utils.to_categorical(test_labels)
+    print("train labels shape: ", train_labels.shape)
+    print("test labels shape: ", test_labels.shape)
+
+    # transform color like: [0, 255] -> 0.xxx
+    print('\n>>> transform image color into float32')
+    train_images = train_images.astype('float32') / 255
+    test_images = test_images.astype('float32') / 255
+
+    # reshape image like: (60000, 28, 28) -> (60000, 28, 28, 1)
+    print('\n>>> reshape image with color channel')
+    train_images = train_images.reshape((60000, 28, 28, 1))
+    test_images = test_images.reshape((10000, 28, 28, 1))
+    print("train images shape: ", train_images.shape)
+    print("test images shape: ", test_images.shape)
+
+    print('\n>>> build model')
+    model = tf.keras.models.Sequential([
+        tf.keras.layers.Conv2D(32, (3, 3), activation=tf.nn.relu, input_shape=(28, 28, 1)),
+        tf.keras.layers.MaxPooling2D((2, 2)),
+        tf.keras.layers.Conv2D(64, (3, 3), activation=tf.nn.relu),
+        tf.keras.layers.MaxPooling2D((2, 2)),
+        tf.keras.layers.Conv2D(64, (3, 3), activation=tf.nn.relu),
+        tf.keras.layers.Flatten(),
+        tf.keras.layers.Dense(64, activation=tf.nn.relu),
+        tf.keras.layers.Dense(10, activation=tf.nn.softmax)
+    ])
+    model.compile(optimizer='rmsprop',
+                  loss='categorical_crossentropy',
+                  metrics=['accuracy'])
+    model.summary()
+
+    print('\n>>> train the model')
+    early_stopping = tf.keras.callbacks.EarlyStopping(
+        monitor='loss', min_delta=0.0005, patience=3, verbose=1, mode='auto',
+        baseline=None, restore_best_weights=True
+    )
+    model.fit(train_images, train_labels, epochs=100, batch_size=64, callbacks=[early_stopping])
+
+    print('\n>>> evaluate the model')
+    test_loss, test_acc = model.evaluate(test_images, test_labels)
+    print("lost: %f, accuracy: %f" % (test_loss, test_acc))
+
+    print('\n>>> save the keras model as %s' % model_name)
+    model.save(model_name)
+
+
+if __name__ == '__main__':
+
+    if not os.path.exists(MODEL_NAME_H5):
+        build_model(MODEL_NAME_H5)
+
+    if not os.path.exists(MODEL_NAME_TFLITE):
+        print('\n>>> save the tflite model as %s' % MODEL_NAME_TFLITE)
+        converter = tf.lite.TFLiteConverter.from_keras_model(tf.keras.models.load_model(MODEL_NAME_H5))
+        tflite_model = converter.convert()
+        with open(MODEL_NAME_TFLITE, "wb") as f:
+            f.write(tflite_model)
+
+    if not os.path.exists(DEFAULT_QUAN_MODEL_NAME_TFLITE):
+        print('\n>>> save the default quantized model as %s' % DEFAULT_QUAN_MODEL_NAME_TFLITE)
+        converter = tf.lite.TFLiteConverter.from_keras_model(tf.keras.models.load_model(MODEL_NAME_H5))
+        converter.optimizations = [tf.lite.Optimize.DEFAULT]
+        tflite_model = converter.convert()
+        with open(DEFAULT_QUAN_MODEL_NAME_TFLITE, "wb") as f:
+            f.write(tflite_model)
+
+    if not os.path.exists(FULL_QUAN_MODEL_NAME_TFLITE):
+        mnist = tf.keras.datasets.mnist
+        (train_images, _), (_, _) = mnist.load_data()
+        train_images = train_images.astype('float32') / 255
+        train_images = train_images.reshape((60000, 28, 28, 1))
+        def representative_data_gen():
+            for input_value in tf.data.Dataset.from_tensor_slices(train_images).batch(1).take(100):
+                yield [input_value]
+        print('\n>>> save the full quantized model as %s' % DEFAULT_QUAN_MODEL_NAME_TFLITE)
+        converter = tf.lite.TFLiteConverter.from_keras_model(tf.keras.models.load_model(MODEL_NAME_H5))
+        converter.optimizations = [tf.lite.Optimize.DEFAULT]
+        converter.representative_dataset = representative_data_gen
+        converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
+        converter.inference_input_type = tf.uint8
+        converter.inference_output_type = tf.uint8
+        tflite_model = converter.convert()
+        with open(FULL_QUAN_MODEL_NAME_TFLITE, "wb") as f:
+            f.write(tflite_model)