若该文为原创文章，转载请注明原文出处。
基于正点原子的ATK-DLRK3568部署测试。
花卉图像分类任务，使用使用 tf.keras.Sequential 模型，简单构建模型，然后转换成 RKNN 模型部署到ATK-DLRK3568板子上。
在 PC 使用 Windows 系统安装 tensorflow，并创建虚拟环境进行训练，然后切换到VM下的RK3568环境，使用rknn-toolkit2把模型转成rknn模型部署到RK3568板子上测试。
一、介绍
TensorFlow 是一个基于数据流编程（dataflow programming）的符号数学系统，被广泛应用于机器学习（machine learning）算法的编程实现，其前身是谷歌的神经网络算法库 DistBelief。
使用 tf.keras.Sequential 模型对花卉图像进行分类。
二、环境搭建
1、创建虚拟环境
conda create -n tensorflow_env python=3.8 -y2、激活环境
conda activate tensorflow_env3、安装环境
pip install numpy
pip install tensorflow
pip install pillow三、训练
1、下载数据集
https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz数据集不好下载，自行处理。
2、训练
tensorflow_classification.py
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.models import Sequential
# 获取
import pathlib
#dataset_url = "https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz"
#data_dir = tf.keras.utils.get_file('flower_photos', origin=dataset_url, untar=True)
data_dir = './flower_photos'
data_dir = pathlib.Path(data_dir)
batch_size = 32
img_height = 180
img_width = 180
# 划分数据
train_ds = tf.keras.utils.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.utils.image_dataset_from_directory(
data_dir,
validation_split=0.2,
subset="validation",
seed=123,
image_size=(img_height, img_width),
batch_size=batch_size)
class_names = train_ds.class_names
#print(class_names)
# 处理数据
normalization_layer = layers.Rescaling(1./255)
train_ds = train_ds.map(lambda x, y: (normalization_layer(x), y))
val_ds = val_ds.map(lambda x, y: (normalization_layer(x), y))
num_classes = len(class_names)
data_augmentation = keras.Sequential(
[
layers.RandomFlip("horizontal",
input_shape=(img_height,
img_width,
3)),
layers.RandomRotation(0.1),
layers.RandomZoom(0.1),
]
)
model = Sequential([
data_augmentation,
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, name="outputs")
])
model.compile(optimizer='adam',
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=['accuracy'])
model.summary()
# 训练模型
epochs=15
history = model.fit(
train_ds,
validation_data=val_ds,
epochs=epochs,
)
# 测试模型
#sunflower_url = "https://storage.googleapis.com/download.tensorflow.org/example_images/592px-Red_sunflower.jpg"
#sunflower_path = tf.keras.utils.get_file('Red_sunflower', origin=sunflower_url)
sunflower_path = './test_180.jpg'
img = tf.keras.utils.load_img(
sunflower_path, target_size=(img_height, img_width)
)
img_array = tf.keras.utils.img_to_array(img)
img_array = tf.expand_dims(img_array, 0) # Create a batch
predictions = model.predict(img_array)
score = tf.nn.softmax(predictions[0])
print(
"This image most likely belongs to {} with a {:.2f} percent confidence."
.format(class_names[np.argmax(score)], 100 * np.max(score))
)
# Convert the model.
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
# Save the model.
with open('model.tflite', 'wb') as f:
f.write(tflite_model)代码有点需要注意，代码屏蔽了下载的功能，所以需要预先下载数据集，如果没有下载数据集，就需要把下载的代码开启。
#dataset_url = "https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz"
#data_dir = tf.keras.utils.get_file('flower_photos', origin=dataset_url, untar=True)执行下面命令开始训练：
python tensorflow_classification.py
等待一会，会生成model.tflite模型文件。
四、RKNN模型转换
转换代码通过下面代码：
rknn_transfer.py
import numpy as np
import cv2
from rknn.api import RKNN
import tensorflow as tf
img_height = 180
img_width = 180
IMG_PATH = 'test.jpg'
class_names = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']
if __name__ == '__main__':
# Create RKNN object
#rknn = RKNN(verbose='Debug')
rknn = RKNN()
# Pre-process config
print('--> Config model')
rknn.config(mean_values=[0, 0, 0], std_values=[255, 255, 255], target_platform='rk3568')
print('done')
# Load model
print('--> Loading model')
ret = rknn.load_tflite(model='model.tflite')
if ret != 0:
print('Load model failed!')
exit(ret)
print('done')
# Build model
print('--> Building model')
ret = rknn.build(do_quantization=False)
#ret = rknn.build(do_quantization=True,dataset='./dataset.txt')
if ret != 0:
print('Build model failed!')
exit(ret)
print('done')
# Export rknn model
print('--> Export rknn model')
ret = rknn.export_rknn('./model.rknn')
if ret != 0:
print('Export rknn model failed!')
exit(ret)
print('done')
#Init runtime environment
print('--> Init runtime environment')
ret = rknn.init_runtime()
# if ret != 0:
# print('Init runtime environment failed!')
# exit(ret)
print('done')
img = cv2.imread(IMG_PATH)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img,(180,180))
img = np.expand_dims(img, 0)
#print('--> Accuracy analysis')
#rknn.accuracy_analysis(inputs=['./test.jpg'])
#print('done')
print('--> Running model')
outputs = rknn.inference(inputs=[img])
print(outputs)
outputs = tf.nn.softmax(outputs)
print(outputs)
print(
"This image most likely belongs to {} with a {:.2f} percent confidence."
.format(class_names[np.argmax(outputs)], 100 * np.max(outputs))
)
#print("图像预测是:", class_names[np.argmax(outputs)])
print('--> done')
rknn.release()
运行后会生成RKNN模型
五、部署
rknnlite_inference.py
import numpy as np
import cv2
from rknnlite.api import RKNNLite
IMG_PATH = 'test.jpg'
RKNN_MODEL = 'model.rknn'
img_height = 180
img_width = 180
class_names = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']
# Create RKNN object
rknn_lite = RKNNLite()
# load RKNN model
print('--> Load RKNN model')
ret = rknn_lite.load_rknn(RKNN_MODEL)
if ret != 0:
print('Load RKNN model failed')
exit(ret)
print('done')
# Init runtime environment
print('--> Init runtime environment')
ret = rknn_lite.init_runtime()
if ret != 0:
print('Init runtime environment failed!')
exit(ret)
print('done')
# load image
img = cv2.imread(IMG_PATH)
img = cv2.resize(img,(180,180))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.expand_dims(img, 0)
# runing model
print('--> Running model')
outputs = rknn_lite.inference(inputs=[img])
print("result: ", outputs)
print(
"This image most likely belongs to {}."
.format(class_names[np.argmax(outputs)])
)
rknn_lite.release()把上面的rknnlite_inference.py和图片拷贝到开发板上，终端运行即可。
如有侵权，或需要完整代码，请及时联系博主。
