mnn模型从训练-转换-预测

之前写过一个文章转换mnn模型

但是没有从头开始，而是直接使用的一个模型，本文想直接从头到尾直接做一下

训练模型的代码如下：

注意：此代码必须tf2.3及其以上才能运行，以下的代码均在tf2.3运行的

import matplotlib.pyplot as plt
import numpy as np
import os
import PIL
import tensorflow as tf

from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.models import Sequential

batch_size = 2
img_height = 180
img_width = 180
num_classes = 5
epochs = 1
data_dir='flower_photos'

train_ds = tf.keras.preprocessing.image_dataset_from_directory(
  data_dir,
  validation_split=0.2,
  subset="training",
  seed=123,
  image_size=(img_height, img_width),
  batch_size=batch_size)

val_ds = tf.keras.preprocessing.image_dataset_from_directory(
  data_dir,
  validation_split=0.2,
  subset="validation",
  seed=123,
  image_size=(img_height, img_width),
  batch_size=batch_size)

AUTOTUNE = tf.data.experimental.AUTOTUNE
train_ds = train_ds.cache().shuffle(1000).prefetch(buffer_size=AUTOTUNE)
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)

data_augmentation = keras.Sequential(
  [
    layers.experimental.preprocessing.RandomFlip("horizontal", 
                                                 input_shape=(img_height, 
                                                                img_width,
                                                              3)),
    layers.experimental.preprocessing.RandomRotation(0.1),
    layers.experimental.preprocessing.RandomZoom(0.1),
  ]
)

model = Sequential([
  data_augmentation,
  layers.experimental.preprocessing.Rescaling(1./255),
  layers.Conv2D(16, 3, padding='same', activation='relu'),
  layers.MaxPooling2D(),
  layers.Conv2D(32, 3, padding='same', activation='relu'),
  layers.MaxPooling2D(),
  layers.Conv2D(64, 3, padding='same', activation='relu'),
  layers.MaxPooling2D(),
  layers.Dropout(0.2),
  layers.Flatten(),
  layers.Dense(128, activation='relu'),
  layers.Dense(num_classes)
])

model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

print(model.summary())
history = model.fit(
  train_ds,
  validation_data=val_ds,
  epochs=epochs
)

#tf.saved_model.save(model,'saved_model')
model.save('flow.h5')

converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
open("flow.tflite", "wb").write(tflite_model)

训练的代码和数据是来自TF官网的案例

训练成功后会保存一个h5和tflite文件

我们可以测试一下h5模型：

import tensorflow as tf
import cv2
import numpy as np

image = cv2.imread('397.jpg')
image=cv2.resize(image,(180,180))
image=image[np.newaxis,:,:,:]
print(image.shape)
loaded_model = tf.keras.models.load_model('flow.h5')
print(loaded_model.predict(image))

运行结果：

[[ 0.18912865  1.8175818  -0.41742945 -1.3345877  -0.3854989 ]]

看一下tflite的预测结果：

import tensorflow as tf
import cv2
import numpy as np

image = cv2.imread('397.jpg')
image=cv2.resize(image,(180,180))
image=image[np.newaxis,:,:,:].astype(np.float32)
print(image.shape)

interpreter = tf.lite.Interpreter(model_path='flow.tflite')
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
interpreter.set_tensor(input_details[0]['index'],image)
interpreter.invoke()
output_data = interpreter.get_tensor(output_details[0]['index'])
print(output_data)

运行结果：

(1, 180, 180, 3)
[[ 0.18912843  1.8175817  -0.41742924 -1.3345877  -0.385499  ]]

两个的结果基本一致

至此证明了tflite模型是正常的

为啥这里直接使用tflite，因为tf现在可以使用tf.lite.TFLiteConverter.from_keras_model直接转keras的模型，而不用费力自己写一大堆代码去转换pb文件了，因为h5没法直接转mnn，要么pb要么tflite才能转，因此经常使用h5转成pb再转mnn

接下来当然需要做的就是tflite转成mnn模型了

mnnconvert -f TFLITE --modelFile flow.tflite --MNNModel model-flow.mnn

运行成功就能得到mnn文件：

 我们可以使用mnn专门的python api直接调用了，暂时本文先不使用c++调用mnn模型文件

import MNN
import cv2
import numpy as np
interpreter = MNN.Interpreter("flows/model-flow.mnn")
session = interpreter.createSession()
input_tensor = interpreter.getSessionInput(session)
image = cv2.imread('397.jpg')
image=cv2.resize(image,(180,180)).astype(np.float32)
#print(type(image[0][9][0]))
tmp_input = MNN.Tensor((1,180, 180,3), MNN.Halide_Type_Float,image,MNN.Tensor_DimensionType_Tensorflow)
#print(tmp_input)
input_tensor.copyFrom(tmp_input)
interpreter.runSession(session)
#print(2)
output_tensor = interpreter.getSessionOutput(session)
output_data=np.array(output_tensor.getData())
print(output_data)

运行成功也能得到上面keras模型和tflite模型的结果

[[ 2.0448089   2.6208436  -0.15651277 -2.080144    0.5771749 ]]

之所以结果数据不一样，因为mnn测试的图片不是一张图

最后附上那张预测图片的图：