|
本帖最后由 jefferyzhang 于 2021-1-9 22:56 编辑
(仅适用于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就可以保存了。
但是因为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()
到这里就可以输出rknn模型用于npu运行了。
|
本帖子中包含更多资源
您需要 登录 才可以下载或查看,没有帐号?立即注册
x
|