| 
 | 
 
 本帖最后由 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
 
 
 
 
 |