[经验]

RK3399Pro入门教程（4）从Tensorflow.Keras到RKNN

mushenmu

2022-3-31 16:23:44

2956

本帖最后由 mushenmu 于 2022-3-31 16:59 编辑

（仅适用于rknn-toolkit < 1.6，后续版本已支持直接从H5转rknn）

tensorflow从1.11开始就推荐使用更高层的Keras API建立模型，代码确实也简洁了不少，但是Keras默认保存的H5格式的模型文件，当前版本的rknn-toolkit(0.9.8)暂时还不支持转换，好在tensorflow还一直保留着他们自己的pb模型格式文件（配置+权重）。

我们直接从Tensorflow官网首页的教程Mnist入手，大概简单说明下如何从Tensorflow.Keras搭建训练模型，然后转RKNN进行使用。

1. import

这是我们要用到的几个模块：

import tensorflow as tf
from tensorflow.python.framework import graph_util

2. 建立Keras模型和训练

我们用Tensorflow的官网首页代码，为了保存模型，需要一点小修改：

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

model = tf.keras.Sequential([
tf.keras.layers.Flatten(input_shape=(28, 28), name="input28x28"),
tf.keras.layers.Dense(512, activation=tf.nn.relu),
tf.keras.layers.Dropout(0.2),
tf.keras.layers.Dense(10, activation=tf.nn.softmax, name="output")
])

model.compile(optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])

model.fit(x_train, y_train, epochs=5) # train

复制代码

其中：

a. 输入层需要明确指定input_shape, 因为是mnist数据，所以他的输入都是28x28大小的单通道灰度图，指定input_shape=(28, 28)，为了后面方便操作，我们显式的给他赋予一个name="input28x28", 注意，赋值给rknn的不是这个名字，而是model.input.op.name打印出来的名字

b. 输出层我们也显式赋予一个name="output"，注意，赋值给rknn的不是这个名字，而是model.output.op.name打印出来的名字

c. model.compile配置了训练细节，然后model.fit就可以直接开始训练了。

3. 保存模型

如果是保存H5格式模型，直接使用model.save就可以保存了。

model.save("./mnist.h5")

复制代码

但是因为RKNN暂时还不支持H5格式，我们需要把它存成tensorflow的pb格式模型

tf.keras.backend.set_learning_phase(0)
session = tf.keras.backend.get_session() # 获取Keras的session

print('input is :', model.input.op.name) # 注意，RKNN传递的input name是这里打印出来的名字
print('output is:', model.output.op.name) # 注意，RKNN传递的output name是这里打印出来的名字

graph = session.graph
with graph.as_default():
output_names = [model.output.op.name]
print('output_names:', output_names)

constant_graph = graph_util.convert_variables_to_constants(session, session.graph_def, output_names) #转换成运算图

with tf.gfile.GFile('./model.pb', mode='wb') as f:
f.write(constant_graph.SerializeToString()) # 写pb文件

也就是多了几行代码而已，还是非常简单的。

4. 用tensorboard或者Netron来查看建立的模型（可选过程）

这时候就可以使用tensorboard来查看这个pb文件的图了，先把上头的constant_graph写成log

writer = tf.summary.FileWriter('./log', constant_graph)
writer.close()

然后用tensorboard命令来查看：

tensorboard --logdir=./log

复制代码

就可以看到刚建立的的tensor图了。

可以看到我们自己定义的输入和输出层名字。这里我们更推荐开源软件Netron来直接查看PB或者H5格式的模型。

5. 转换成RKNN

然后就可以通过RKNN-toolkit转换啦，当然这里要先准备一张训练的图，把图的路径放入dataset.txt文件里，用于rknn量化（定点转换、优化等过程）

# coding=utf-8
from rknn.api import RKNN

if __name__ == '__main__':

# Create RKNN object
rknn = RKNN(verbose=True)

# pre-process config
print('--> config model')
rknn.config(channel_mean_value='0 0 0 255')  # 因为是灰度图，不需要设置reorder_channel进行通道转换
print('done')

# Load tensorflow model
print('--> Loading model')
ret = rknn.load_tensorflow(
      tf_pb='./model.pb',
      inputs=['input28x28_input'],  # 注意，这里的input名字来自于模型转换时候打印出来的mode.input.op.name
      outputs=['output/Softmax'], # 注意，这里的output名字来自于模型转换时候打印出来的mode.output.op.name
      input_size_list=[[28, 28]])
if ret != 0:
      print('Load mtcnn failed! Ret = {}'.format(ret))
      exit(ret)
print('done')

# Build model
print('--> Building model')
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')  # 保存成rknn模型文件
if ret != 0:
      print('Export rknn failed!')
      exit(ret)
print('done')

rknn.release()