cs249r 嵌入式机器学习

#include <TensorFlowLite_ESP32.h>


#include "tensorflow/lite/micro/all_ops_resolver.h"
#include "tensorflow/lite/micro/micro_error_reporter.h"
#include "tensorflow/lite/micro/micro_interpreter.h"
#include "tensorflow/lite/micro/system_setup.h"
#include "tensorflow/lite/schema/schema_generated.h"

#include "main_functions.h"
#include "model.h"
#include "constants.h"
#include "output_handler.h"

// Globals, used for compatibility with Arduino-style sketches.
namespace {
tflite::ErrorReporter* error_reporter = nullptr;
const tflite::Model* model = nullptr;
tflite::MicroInterpreter* interpreter = nullptr;
TfLiteTensor* input = nullptr;
TfLiteTensor* output = nullptr;
int inference_count = 0;

constexpr int kTensorArenaSize = 2000;
uint8_t tensor_arena[kTensorArenaSize];
}  // namespace

// The name of this function is important for Arduino compatibility.
void setup() {

  // Set up logging. Google style is to avoid globals or statics because of
  // lifetime uncertainty, but since this has a trivial destructor it's okay.
  // NOLINTNEXTLINE(runtime-global-variables)
  static tflite::MicroErrorReporter micro_error_reporter;
  error_reporter = &micro_error_reporter;

  // Map the model into a usable data structure. This doesn't involve any
  // copying or parsing, it's a very lightweight operation.
  model = tflite::GetModel(g_model);
  if (model->version() != TFLITE_SCHEMA_VERSION) {
    TF_LITE_REPORT_ERROR(error_reporter,
                         "Model provided is schema version %d not equal "
                         "to supported version %d.",
                         model->version(), TFLITE_SCHEMA_VERSION);
    return;
  }

  // This pulls in all the operation implementations we need.
  // NOLINTNEXTLINE(runtime-global-variables)
  static tflite::AllOpsResolver resolver;

  // Build an interpreter to run the model with.
  static tflite::MicroInterpreter static_interpreter(
      model, resolver, tensor_arena, kTensorArenaSize, error_reporter);
  interpreter = &static_interpreter;

  // Allocate memory from the tensor_arena for the model's tensors.
  TfLiteStatus allocate_status = interpreter->AllocateTensors();
  if (allocate_status != kTfLiteOk) {
    TF_LITE_REPORT_ERROR(error_reporter, "AllocateTensors() failed");
    return;
  }

  // Obtain pointers to the model's input and output tensors.
  input = interpreter->input(0);
  output = interpreter->output(0);

  // Keep track of how many inferences we have performed.
  inference_count = 0;
}

// The name of this function is important for Arduino compatibility.
void loop() {
  // Calculate an x value to feed into the model. We compare the current
  // inference_count to the number of inferences per cycle to determine
  // our position within the range of possible x values the model was
  // trained on, and use this to calculate a value.
  float position = static_cast<float>(inference_count) /
                   static_cast<float>(kInferencesPerCycle);
  float x = position * kXrange;

  // Quantize the input from floating-point to integer
  int8_t x_quantized = x / input->params.scale + input->params.zero_point;
  // Place the quantized input in the model's input tensor
  input->data.int8[0] = x_quantized;

  // Run inference, and report any error
  TfLiteStatus invoke_status = interpreter->Invoke();
  if (invoke_status != kTfLiteOk) {
    TF_LITE_REPORT_ERROR(error_reporter, "Invoke failed on x: %f\n",
                         static_cast<double>(x));
    return;
  }

  // Obtain the quantized output from model's output tensor
  int8_t y_quantized = output->data.int8[0];
  // Dequantize the output from integer to floating-point
  float y = (y_quantized - output->params.zero_point) * output->params.scale;

  // Output the results. A custom HandleOutput function can be implemented
  // for each supported hardware target.
  HandleOutput(error_reporter, x, y);

  // Increment the inference_counter, and reset it if we have reached
  // the total number per cycle
  inference_count += 1;
  if (inference_count >= kInferencesPerCycle) inference_count = 0;
}

//outside the function
tflite::MicroInterpreter* interpreter = nullptr;
//setup()
static tflite::MicroInterpreter static_interpreter(model, resolver, tensor_arena, kTensorArenaSize, error_reporter);
interpreter = &static_interpreter;

def load_model_from_file(model_filename):
	with open(model_filename, "rb") as file:
		buffer_data = file.read()
	model_obj = Model.Model.GetRootAsModel(buffer_data, 0)
	model = Model.ModelT.InitFromObj(model_obj)
	return model

def save_model_to_file(model, model_filename):
	builder = flatbuffers.Builder(1024)
	model_offset = model.Pack(builder)
    builder.Finish(model_offset, file_identifier=b'TFL3')
    model_data = builder.Output()
    with open(model_filename, 'wb') as out_file:
        out_file.write(model_data)
model = load_model_from_file('model.tflite')
for buffer in model.buffers:
	if buffer.data is not None and len(buffer.data) > 1024:
		original_weights = np.frombuffer(buffer.data，dtype=np.float32)
        munged_weights = np .round (original_weights * (1/0.02)) *0.02
    save_model_to_file(model,'model_modified.tflite ')
    
# https://colab.research.google.com/github/tinyMLx/colabs/blob/master/4-4-8-Flatbuffers.ipynb

constexpr int kTensorArenaSize = 2000;
uint8_t tensor_arena[kTensorArenaSize];

!apt-get update && apt-get -qq install xxd
MODEL_TFLITE = 'KWS_yes_no.tflite'
MODEL_TFLITE_MICRO = 'KWS_yes_no.cc'
!xxd -i {MODEL_TFLITE} > {MODEL_TFLITE_MICRO}
REPLACE_TEXT = MODEL_TFLITE.replace('/', '_').replace('.', '_')
!sed -i 's/'{REPLACE_TEXT}'/g_model/g' {MODEL_TFLITE_MICRO}

import tensorflow as tf
print("TensorFlow version:", tf.__version__)

from tensorflow.keras.layers import Dense, Flatten, Conv2D, AveragePooling2D
from tensorflow.keras import Model
from tqdm import tqdm

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

# Add a channels dimension
x_train = x_train[..., tf.newaxis].astype("float32")
x_test = x_test[..., tf.newaxis].astype("float32")

train_ds = tf.data.Dataset.from_tensor_slices(
    (x_train, y_train)).shuffle(10000).batch(32)

test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)

class MyModel(Model):
  def __init__(self):
    super(MyModel, self).__init__()
    self.conv1 = Conv2D(3, 5, activation='relu')
    self.avgpool1 = AveragePooling2D()
    self.flatten = Flatten()
    self.d2 = Dense(10)

  def call(self, x):
    x = self.conv1(x)
    x = self.avgpool1(x)
    x = self.flatten(x)
    return self.d2(x)

# Create an instance of the model
model = MyModel()

loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

optimizer = tf.keras.optimizers.Adam()

train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')

test_loss = tf.keras.metrics.Mean(name='test_loss')
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')

@tf.function
def train_step(images, labels):
  with tf.GradientTape() as tape:
    # training=True is only needed if there are layers with different
    # behavior during training versus inference (e.g. Dropout).
    predictions = model(images, training=True)
    loss = loss_object(labels, predictions)
  gradients = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))

  train_loss(loss)
  train_accuracy(labels, predictions)
    
@tf.function
def test_step(images, labels):
  # training=False is only needed if there are layers with different
  # behavior during training versus inference (e.g. Dropout).
  predictions = model(images, training=False)
  t_loss = loss_object(labels, predictions)

  test_loss(t_loss)
  test_accuracy(labels, predictions)
    
EPOCHS = 50

for epoch in tqdm(range(EPOCHS)):
  # Reset the metrics at the start of the next epoch
  train_loss.reset_states()
  train_accuracy.reset_states()
  test_loss.reset_states()
  test_accuracy.reset_states()

  for images, labels in train_ds:
    train_step(images, labels)

  for test_images, test_labels in test_ds:
    test_step(test_images, test_labels)

  print(
    f'Epoch {epoch + 1}, '
    f'Loss: {train_loss.result()}, '
    f'Accuracy: {train_accuracy.result() * 100}, '
    f'Test Loss: {test_loss.result()}, '
    f'Test Accuracy: {test_accuracy.result() * 100}'
  )

model.save(MODEL_SAVE_DIR) #模型存储

converter = tf.lite.TFLiteConverter.from_saved_model(MODEL_SAVE_DIR) #从模型设置一个converter对象
model_no_quant_tflite = converter.convert() #直接进行convert操作得到一个未经量化的模型，占用较大

open(FLOAT_LITE_MODEL_DIR,'wb').write(model_no_quant_tflite)

# 接下来的模型则是经过优化，关于模型的优化程度可以参考后面的图，这里使用的是MCU平台，因此需要将前面的操作尽可能的转化为INT8。具体的量化设置可以参考这里：https://www.tensorflow.org/lite/performance/post_training_quantization
converter.optimizations = [tf.lite.Optimize.DEFAULT] 
#进行全INT8型转换时必须定义一个representative_dataset来校准整体的范围来保证数据量化正确。（这里并没有明确说明tf是如何进行量化的）https://www.tensorflow.org/lite/performance/quantization_spec
def representative_dataset():
  for i in x_test:
    t = tf.expand_dims(i,0)
    # print(t.shape)
    yield([t])
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8
# Provide a representative dataset to ensure we quantize correctly.
converter.representative_dataset = representative_dataset
model_tflite = converter.convert()

# Save the model to disk
open(QUANT_LITE_MODEL_DIR, "wb").write(model_tflite)


import os


def get_dir_size(dir):
  size = 0
  for f in os.scandir(dir):
    if f.is_file():
      size += f.stat().st_size
    elif f.is_dir():
      size += get_dir_size(f.path)
  return size

# Calculate size
size_tf = get_dir_size(MODEL_SAVE_DIR)
size_no_quant_tflite = os.path.getsize(FLOAT_LITE_MODEL_DIR)
size_tflite = os.path.getsize(QUANT_LITE_MODEL_DIR)

print("original size:{}\nno quant size:{}\nquant size:{}".format(size_tf,size_no_quant_tflite,size_tflite))

import numpy as np
# 测试函数，用于测试模型在压缩后的准确程度
def predict_tflite(tflite_model, img):
  img_array = tf.expand_dims(img, 0)

  # Initialize the TFLite interpreter
  interpreter = tf.lite.Interpreter(model_content=tflite_model)
  interpreter.allocate_tensors()

  input_details = interpreter.get_input_details()[0]
  output_details = interpreter.get_output_details()[0]

  # If required, quantize the input layer (from float to integer)
  # 这里的必须要通过函数来获得其中的参数来计算输入和推理的真值：real_value = (int8_value - zero_point) *scale
  input_scale, input_zero_point = input_details["quantization"]
  if (input_scale, input_zero_point) != (0.0, 0):
    img_array = np.multiply(img_array, 1.0 / input_scale) + input_zero_point
    img_array = img_array.astype(input_details["dtype"])

  # Invoke the interpreter
  interpreter.set_tensor(input_details["index"], img_array)
  interpreter.invoke()
  pred = interpreter.get_tensor(output_details["index"])[0]

  # If required, dequantized the output layer (from integer to float)
  output_scale, output_zero_point = output_details["quantization"]
  if (output_scale, output_zero_point) != (0.0, 0):
    pred = pred.astype(np.float32)
    pred = np.multiply((pred - output_zero_point), output_scale)

  predicted_label_index = np.argmax(pred)
  predicted_score = pred[predicted_label_index]
  return (predicted_label_index, predicted_score)
#%%
def run_tflite_test(model_file):
  correct_count = 0
  wrong_count = 0
  discarded_count = 0
  for i,j in enumerate(x_test):
      index, score = predict_tflite(model_file, x_test[i])
      if score < 0.75:
        discarded_count += 1
        continue
      if index == y_test[i]:
        correct_count += 1
      else:
        wrong_count += 1
        print("[%s] expected, [%s] found with score [%f]" % (y_test[i], index, score))

  correct_percentage = (correct_count / (correct_count + wrong_count)) * 100


  print("%.1f%% correct (N=%d, %d unknown)" % (correct_percentage, (correct_count + wrong_count), discarded_count))
#%%
run_tflite_test(model_no_quant_tflite)
#%%
run_tflite_test(model_tflite)

int8_t getData(signed char *input)
{
    if (WiFi.status() == WL_CONNECTED)
    {
        HTTPClient http;
        http.begin(serverPath);
        int httpCode = http.GET();

        char *buf = (char *)malloc(4096 * sizeof(char));
        String s = http.getString();
        s.toCharArray(buf, 4096);

        std::stringstream ss(buf);
        for (int i = 0; i < 784; i++)
        {
            int t;
            ss >> t;
            input[i] = t;
        }
        int8_t label = 0;
        while(ss)
        ss >> label;

        return label-48;
    }
    else
    {
        return -1;
    }
}

void loop() {
    int8_t label = getData(model_input->data.int8);
    if(label == -1)
        error_reporter->Report("failed to get vilid data!");
    TfLiteStatus invoke_status = interpreter->Invoke();
    if (invoke_status != kTfLiteOk) {
    error_reporter->Report("Invoke failed.\n");
    return;
    }
    int8_t* output = interpreter->output(0)->data.int8;
    int argmax=0;
    for(int j=0;j<10;j++)
    {
        if (output[j] > output[argmax])
            argmax = j;
    }
    Serial.printf("predicted label:%d\t true value:%d\n",argmax,label);
    delay(1000);
}