之前写过一个文章转换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测试的图片不是一张图
最后附上那张预测图片的图: