Toybrick

RK3399Pro入门教程(4)从Tensorflow.Keras到RKNN

jefferyzhang

超级版主

积分
1840
发表于 2019-2-18 16:42:15    查看: 4800|回复: 28 | [复制链接]    打印 | 显示全部楼层
本帖最后由 jefferyzhang 于 2019-8-26 21:21 编辑


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

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


1. import


这是我们要用到的几个模块:
  1. import tensorflow as tf
  2. from tensorflow.python.framework import graph_util
复制代码


2. 建立Keras模型和训练


我们用Tensorflow的官网首页代码,为了保存模型,需要一点小修改:
  1. mnist = tf.keras.datasets.mnist
  2.     (x_train, y_train), (x_test, y_test) = mnist.load_data()
  3.     x_train, x_test = x_train / 255.0, x_test / 255.0

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

  10.     model.compile(optimizer='adam',
  11.                   loss='sparse_categorical_crossentropy',
  12.                   metrics=['accuracy'])
复制代码
  1. 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就可以保存了。
  1. model.save("./mnist.h5")
复制代码


但是因为RKNN暂时还不支持H5格式,我们需要把它存成tensorflow的pb格式模型
  1. tf.keras.backend.set_learning_phase(0)
  2. session = tf.keras.backend.get_session()   # 获取Keras的session

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

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

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

  10. with tf.gfile.GFile('./model.pb', mode='wb') as f:
  11.     f.write(constant_graph.SerializeToString())   # 写pb文件
复制代码
也就是多了几行代码而已,还是非常简单的。


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


这时候就可以使用tensorboard来查看这个pb文件的图了,先把上头的constant_graph写成log
  1. writer = tf.summary.FileWriter('./log', constant_graph)
  2.         writer.close()
复制代码
然后用tensorboard命令来查看:
  1. >  tensorboard --logdir=./log
复制代码
就可以看到刚建立的的tensor图了。


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


5. 转换成RKNN



然后就可以通过RKNN-toolkit转换啦,当然这里要先准备一张训练的图,把图的路径放入dataset.txt文件里,用于rknn量化(定点转换、优化等过程)
  1. # coding=utf-8
  2. from rknn.api import RKNN

  3. if __name__ == '__main__':

  4.     # Create RKNN object
  5.     rknn = RKNN(verbose=True)

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

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

  21.     # Build model
  22.     print('--> Building model')
  23.     ret = rknn.build(do_quantization=True, dataset='./dataset.txt')  # 量化模型
  24.     if ret != 0:
  25.         print('Build model failed!')
  26.         exit(ret)
  27.     print('done')

  28.     # Export rknn model
  29.     print('--> Export RKNN model')
  30.     ret = rknn.export_rknn('./model.rknn')  # 保存成rknn模型文件
  31.     if ret != 0:
  32.         print('Export rknn failed!')
  33.         exit(ret)
  34.     print('done')


  35.     rknn.release()
复制代码

到这里就可以输出rknn模型用于npu运行了。

本帖子中包含更多资源

您需要 登录 才可以下载或查看,没有帐号?立即注册

x
回复

使用道具 举报

samtu

中级会员

积分
303
发表于 2019-5-4 13:41:48 | 显示全部楼层
我想问一下,这个模型的训练也可以在3399PRO上面来做吗,会不会太慢,
回复

使用道具 举报

jefferyzhang

超级版主

积分
1840
 楼主| 发表于 2019-5-5 18:05:49 | 显示全部楼层
samtu 发表于 2019-5-4 13:41
我想问一下,这个模型的训练也可以在3399PRO上面来做吗,会不会太慢,

不能的 ,我们4块titanV跑了一晚上收敛都不是特别理想,在3399上用cpu做那基本不可能。。。
回复

使用道具 举报

keyman_sysu

注册会员

积分
146
发表于 2019-5-14 12:33:33 | 显示全部楼层
我发现就算输出是7*7*128的digits,RKNN输出也是一维的6272的digits。那么channel_first是否意味着我要先reshape成为128*7*7,然后再transpose成为7*7*128啊?
回复

使用道具 举报

samtu

中级会员

积分
303
发表于 2019-5-14 16:34:51 | 显示全部楼层
jefferyzhang 发表于 2019-5-5 18:05
不能的 ,我们4块titanV跑了一晚上收敛都不是特别理想,在3399上用cpu做那基本不可能。。。 ...

楼主,我还有一个问题,为什么要用RKNN来转换呢,难道直接用TENSORFLOW 不行吗,还是转换后才可以用到你们的NPU,
回复

使用道具 举报

samtu

中级会员

积分
303
发表于 2019-5-14 21:06:12 | 显示全部楼层
4块titanV,10万块,利器.,
回复

使用道具 举报

jefferyzhang

超级版主

积分
1840
 楼主| 发表于 2019-5-15 10:42:12 | 显示全部楼层
samtu 发表于 2019-5-14 21:06
4块titanV,10万块,利器.,

近些年更好的人脸检测论文应该是arcface, 收敛和分类应该是优秀的
回复

使用道具 举报

jefferyzhang

超级版主

积分
1840
 楼主| 发表于 2019-5-15 10:50:40 | 显示全部楼层
samtu 发表于 2019-5-14 16:34
楼主,我还有一个问题,为什么要用RKNN来转换呢,难道直接用TENSORFLOW 不行吗,还是转换后才可以用到你 ...

NPU只能读取rknn格式模型,就跟tf只能读取pb格式一样。需要做一次转换。
回复

使用道具 举报

samtu

中级会员

积分
303
发表于 2019-5-20 22:31:44 | 显示全部楼层
你们自己训练是从头开始吗,还是拿别人的模型(如VGG,RESNET。。。,拆开,再从FC层后面加分类什么之 类的?自己训练,没有什么数据呢,还要自己一个个的标记,
回复

使用道具 举报

jefferyzhang

超级版主

积分
1840
 楼主| 发表于 2019-5-21 09:52:45 | 显示全部楼层
samtu 发表于 2019-5-20 22:31
你们自己训练是从头开始吗,还是拿别人的模型(如VGG,RESNET。。。,拆开,再从FC层后面加分类什么之 类的 ...

当然用开源数据集。RK发布的模型有的是rk自己训练,有的是下载预训练模型。
对于产品开发来说,大部分的模型需要重新自己训练
回复

使用道具 举报

您需要登录后才可以回帖 登录 | 立即注册

本版积分规则

产品中心 购买渠道 开源社区 Wiki教程 资料下载 关于Toybrick


快速回复 返回顶部 返回列表