Toybrick

Toybrick»开源社区 › 技术论坛 › 人工智能 › RK3399Pro入门教程（9）MNIST RKNN量化教程

RK3399Pro入门教程（9）MNIST RKNN量化教程

peng

注册会员

积分
169
楼主
发表于 2019-9-20 08:34:19    查看: 43355|回复: 11 | [复制链接]    打印 | 显示全部楼层

一.搭建网络
    model.py,同tensorflow官网mnist的例子差不多，不懂的可以去官网看下官网代码解析
  1. import tensorflow as tf

  3. #这里输入采用28x28方便之后进行rknn量化
  4. x = tf.placeholder("float", [None, 28,28],name='x')
  5. y_ = tf.placeholder("float", [None,10],name='y_')
  6. keep_prob = tf.placeholder("float", name='keep_prob')
  7. def weight_variable(shape,name):
  8.     initial = tf.truncated_normal(shape, stddev=0.1)
  9.     return tf.Variable(initial,name=name)

  11. def bias_variable(shape,name):
  12.     initial = tf.constant(0.1, shape=shape)
  13.     return tf.Variable(initial,name=name)

  15. def conv2d(x, W):
  16.     return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')

  18. def max_pool_2x2(x):
  19.     return tf.nn.max_pool(x, ksize=[1, 2, 2, 1],
  20.                         strides=[1, 2, 2, 1], padding='SAME')

  22. # convolution layer
  23. def lenet5_layer(input, weight, bias,weight_name,bias_name):
  24.     W_conv = weight_variable(weight,weight_name)
  25.     b_conv = bias_variable(bias,bias_name)
  26.     h_conv = tf.nn.relu(conv2d(input, W_conv) + b_conv)
  27.     return max_pool_2x2(h_conv)
  28. # connected layer
  29. def dense_layer(layer, weight, bias,weight_name,bias_name):
  30.     W_fc = weight_variable(weight,weight_name)
  31.     b_fc = bias_variable(bias,bias_name)
  32.     return tf.nn.relu(tf.matmul(layer, W_fc) + b_fc)

  34. def build_model(is_training):
  35.     #first conv
  36.     x_image = tf.reshape(x, [-1,28,28,1])
  37.     W_conv1 = [5, 5, 1, 32]
  38.     b_conv1 = [32]
  39.     layer1 = lenet5_layer(x_image,W_conv1,b_conv1,'W_conv1','b_conv1')
  40.     #second conv
  41.     W_conv2 = [5, 5, 32, 64]
  42.     b_conv2 = [64]
  43.     layer2 = lenet5_layer(layer1,W_conv2,b_conv2,'W_conv2','b_conv2')
  44.     #third conv
  45.     W_fc1 = [7 * 7 * 64, 1024]
  46.     b_fc1 = [1024]
  47.     layer2_flat = tf.reshape(layer2, [-1, 7*7*64])
  48.     layer3 = dense_layer(layer2_flat,W_fc1,b_fc1,'W_fc1','b_fc1')
  49.     #softmax
  50.     W_fc2 = weight_variable([1024, 10],'W_fc2')
  51.     b_fc2 = bias_variable([10],'b_fc2')
  52.     if is_training:
  53.         #dropout
  54.         h_fc1_drop = tf.nn.dropout(layer3, keep_prob)
  55.         finaloutput=tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2,name="y_conv")
  56.     else:
  57.         finaloutput=tf.nn.softmax(tf.matmul(layer3, W_fc2) + b_fc2,name="y_conv")
  58.     print('finaloutput:', finaloutput)
  59.     return finaloutput
复制代码
二.训练网络
    train.py,这里代码兼容了tf伪量化的代码，这里我们把create_training_graph()传入的参数is_quantify设为False就可以了，由于mnist拿到的train和test数据shape都是（784，),这里定义了一个reshape_batch函数把train时的batch以及test时的输入都reshape成（28,28），具体代码如下：
  1. # -*- coding=utf-8 -*-
  2. import tensorflow as tf
  3. from tensorflow.examples.tutorials.mnist import input_data
  4. from model import build_model, x, y_, keep_prob

  6. mnist = input_data.read_data_sets("../MNIST_data/", one_hot=True)
  7. def create_training_graph(is_quantify):
  8.     #创建训练图，加入create_training_graph：
  9.     g = tf.get_default_graph()   # 给create_training_graph的参数，默认图
  10.     #调用网络定义，也就是拿到输出
  11.     y_conv = build_model(is_training=True)    #这里的is_training设置为True，因为前面模型定义写了训练时要用到dropout
  12.     #损失函数
  13.     cross_entropy = -tf.reduce_sum(y_*tf.log(y_conv))
  14.     print('cost:', cross_entropy)
  15.     if is_quantify:
  16.         # 加入 create_training_graph函数，注意位置要在loss之后， optimize之前
  17.         tf.contrib.quantize.create_training_graph(input_graph=g, quant_delay=0)
  18.     #  optimize
  19.     optimize = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy)
  20.     #计算准确率
  21.     correct_prediction = tf.equal(tf.argmax(y_conv,1), tf.argmax(y_,1))
  22.     # 给出识别准确率[这会返回我们一个布尔值的列表.为了确定哪些部分是正确的，我们要把它转换成浮点值，然后再示均值。 比如, [True, False, True, True] 会转换成 [1,0,1,1] ，从而它的准确率就是0.75.]   
  23.     accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"))

  25.     # 返回所需数据，供训练使用
  26.     return dict(
  27.         x=x,
  28.         y=y_,
  29.         keep_prob=keep_prob,
  30.         optimize=optimize,
  31.         cost=cross_entropy,
  32.         correct_prediction=correct_prediction,
  33.         accuracy=accuracy,
  34.     )

  36. def reshape_batch(batch):
  37.     rebatch = []
  38.     for item in batch:
  39.         b = item.reshape(28,28)
  40.         rebatch.append(b)
  41.     return rebatch
  42. #开始训练
  43. def train_network(graph,ckpt,point_dir,pbtxt):
  44.     # 初始化
  45.     init = tf.global_variables_initializer()
  46.     # 调用Saver函数保存所需文件
  47.     saver = tf.train.Saver()
  48.     # 创建上下文，开始训练sess.run(init)
  49.     with tf.Session() as sess:
  50.         sess.run(init)
  51.         # 一共训练两万次，准确率达到百分99以上
  52.         for i in range(20000):
  53.         # 每次处理50张图片
  54.             batch = mnist.train.next_batch(50)
  55.             # 每100次保存并打印一次准确率等
  56.             if i % 100 == 0:
  57.             # feed_dict喂数据，数据全reshape成28x28
  58.                 train_accuracy = sess.run([graph['accuracy']], feed_dict={
  59.                                                                            graph['x']:reshape_batch(batch[0]),    # batch[0]存的图片数据
  60.                                                                            graph['y']:batch[1],    # batch[1]存的标签
  61.                                                                            graph['keep_prob']: 1.0})
  62.                 print("step %d, training accuracy %g"%(i, train_accuracy[0]))
  63.             sess.run([graph['optimize']], feed_dict={
  64.                                                        graph['x']:reshape_batch(batch[0]),
  65.                                                        graph['y']:batch[1],
  66.                                                        graph['keep_prob']:0.5})
  67.         test_accuracy = sess.run([graph['accuracy']], feed_dict={
  68.                                                                   graph['x']: reshape_batch(mnist.test.images),
  69.                                                                   graph['y']: mnist.test.labels,
  70.                                                                   graph['keep_prob']: 1.0})
  71.         print("Test accuracy %g" % test_accuracy[0])
  72.         # 保存ckpt(checkpoint)和pbtxt。记得把路径改成自己的路径
  73.         saver.save(sess, ckpt)
  74.         tf.train.write_graph(sess.graph_def,point_dir,pbtxt, True)
  75.         print(tf.trainable_variables())
  76.         print(tf.get_variable('W_fc2',[1024, 10]).value)


  79. if __name__ == "__main__":
  80.     ckpt = './checkpoint/mnist.ckpt'
  81.     point_dir = './checkpoint'
  82.     pbtxt = 'mnist.pbtxt'
  83.     g1 = create_training_graph(False)
  84.     train_network(g1,ckpt,point_dir,pbtxt)

复制代码三.保存网络参数
    freese.py,将网络中参数和变量从ckpt中读出来，保存为pb文件，同上一步一样，将frozen函数的is_quantify设为False就可以了：
  1. import tensorflow as tf
  2. import os.path
  3. from model import build_model
  4. from tensorflow.python.framework import graph_util

  6. # 创建推理图
  7. def create_inference_graph():
  8.     """Build the mnist model for evaluation."""
  9. # 调用网络，Create an output to use for inference.
  10.     logits = build_model(is_training=False)
  11.     return logits
  12.     # # 得到分类输出  
  13.     # tf.nn.softmax(logits, name='output')
  14. def load_variables_from_checkpoint(sess, start_checkpoint):
  15.     """Utility function to centralize checkpoint restoration.
  16.     Args:
  17.       sess: TensorFlow session.
  18.       start_checkpoint: Path to saved checkpoint on disk.
  19.     """
  20.     saver = tf.train.Saver(tf.global_variables())
  21.     saver.restore(sess, start_checkpoint)

  23. def frozen(is_quantify,ckpt,pbtxt):
  24.     # Create the model and load its weights.
  25.     init = tf.global_variables_initializer()
  26.     with tf.Session() as sess:
  27.         sess.run(init)
  28. # 推理图
  29.         logits = create_inference_graph()  
  30. # 加入create_eval_graph()，转化为tflite可接受的格式。以下语句中有路径的，记得改路径。
  31.         if is_quantify:
  32.             tf.contrib.quantize.create_eval_graph()
  33.         load_variables_from_checkpoint(sess, ckpt)
  34.         # Turn all the variables into inline constants inside the graph and save it.
  35. # 固化 frozen：ckpt + pbtxt
  36.         frozen_graph_def = graph_util.convert_variables_to_constants(
  37.             sess, sess.graph_def, ['y_conv'])
  38. # 保存最终的pb模型
  39.         tf.train.write_graph(
  40.             frozen_graph_def,
  41.             os.path.dirname(pbtxt),
  42.             os.path.basename(pbtxt),
  43.             as_text=False)
  44.         tf.logging.info('Saved frozen graph to %s', pbtxt)

  46. if __name__ == "__main__":
  47.     ckpt = './checkpoint/mnist.ckpt'
  48.     pbtxt = 'mnist_frozen_graph.pb'
  49.     frozen(False,ckpt,pbtxt)
  50.     #is_quantify False   mnist_frozen_graph_not_28x28.pb
  51.     # ckpt = './checkpoint_not/mnist.ckpt'
  52.     # pbtxt = 'mnist_frozen_graph_not.pb'
  53.     # frozen(False,ckpt,pbtxt)
  54.     # ckpt = './test/mnist.ckpt'
  55.     # pbtxt = 'test.pb'
  56.     # frozen(False,ckpt,pbtxt)
复制代码四.将pb模型转为rknn
     由于量化rknn模型需要相应图片集，因此我们先要获取相应的数据集进入mnist数据目录下，解压t10k-images-idx3-ubyte.gz，然后运行get_image.py,将原先压缩的数据转为图片，同时得到量化需要的dataset.txt文件。
    get_image.py：
  1. import struct
  2. import numpy as np
  3. #import matplotlib.pyplot as plt
  4. import PIL.Image
  5. from PIL import Image
  6. import os

  8. os.system("mkdir ../MNIST_data/mnist_test")
  9. filename='../MNIST_data/t10k-images.idx3-ubyte'
  10. dataset = './dataset.txt'
  11. binfile=open(filename,'rb')
  12. buf=binfile.read()
  13. index=0
  14. data_list = []
  15. magic,numImages,numRows,numColumns=struct.unpack_from('>IIII',buf,index)
  16. index+=struct.calcsize('>IIII')
  17. for image in range(0,numImages):
  18.     im=struct.unpack_from('>784B',buf,index)
  19.     index+=struct.calcsize('>784B')
  20.     im=np.array(im,dtype='uint8')
  21.     im=im.reshape(28,28)
  22.     im=Image.fromarray(im)
  23.     im.save('../MNIST_data/mnist_test/test_%s.jpg'%image,'jpeg')
  24.     data_list.append('../MNIST_data/mnist_test/test_%s.jpg\n'%image)
  25. with open(dataset,'w+') as ff:
  26.     ff.writelines(data_list)
复制代码rknn_transfer.py:
  1. from rknn.api import RKNN

  3. def common_transfer(pb_name,export_name):
  4.         ret = 0
  5.         #看具体log 传入verbose=True
  6.         rknn = RKNN()
  7.         #灰度图无需此步操作
  8.         # rknn.config(channel_mean_value='', reorder_channel='')
  9.         print('--> Loading model')

  11.         ret = rknn.load_tensorflow(
  12.                 tf_pb='./mnist_frozen_graph.pb',
  13.                 inputs=['x'],
  14.                 outputs=['y_conv'],
  15.                 input_size_list=[[28,28,1]])
  16.         if ret != 0:
  17.                 print('load_tensorflow error')
  18.                 rknn.release()
  19.                 return ret
  20.         print('done')
  21.         print('--> Building model')
  22.         rknn.build(do_quantization=False)
  23.         print('done')
  24.         # 导出保存rknn模型文件
  25.         rknn.export_rknn('./mnist.rknn')
  26.         # Release RKNN Context
  27.         rknn.release()
  28.         return ret

  30. def quantify_transfer(pb_name,dataset_name,export_name):
  31.         ret = 0
  32.         print(pb_name,dataset_name,export_name)
  33.         rknn = RKNN()
  34.         rknn.config(channel_mean_value='', reorder_channel='',quantized_dtype='dynamic_fixed_point-8')
  35.         print('--> Loading model')
  36.         ret = rknn.load_tensorflow(
  37.                 tf_pb=pb_name,
  38.                 inputs=['x'],
  39.                 outputs=['y_conv'],
  40.                 input_size_list=[[28,28,1]])
  41.         if ret != 0:
  42.                 print('load_tensorflow error')
  43.                 rknn.release()
  44.                 return ret
  45.         print('done')
  46.         print('--> Building model')
  47.         rknn.build(do_quantization=True,dataset=dataset_name)
  48.         print('done')
  49.         # 导出保存rknn模型文件
  50.         rknn.export_rknn(export_name)
  51.         # Release RKNN Context
  52.         rknn.release()
  53.         return ret
  54. if __name__ == '__main__':
  55.         #pb转化为rknn模型
  56.         pb_name = './mnist_frozen_graph.pb'
  57.         export_name = './mnist.rknn'
  58.         ret = common_transfer(pb_name,export_name)
  59.         if ret != 0:
  60.                 print('======common transfer error !!===========')
  61.         else:
  62.                 print('======common transfer ok !!===========')
  63.         dataset_name = './dataset.txt'
  64.         export_name = './mnist_quantization.rknn'
  65.         #pb转化为量化的rknn模型
  66.         quantify_transfer(pb_name,dataset_name,export_name)
  67.         if ret != 0:
  68.                 print('======quantization transfer 10000 error !!===========')
  69.         else:
  70.                 print('======quantization transfer 10000 ok !!===========')

复制代码五.对比pb和rknn的推理结果，比较他们的准确度
     分别运行tf_predict.py,rknn_predict.py得到tf模型,rknn模型,量化的rknn模型的运行结果:
     tf_predict.py
  1. #! -*- coding: utf-8 -*-
  2. from __future__ import absolute_import, unicode_literals
  3. from tensorflow.examples.tutorials.mnist import input_data
  4. import tensorflow as tf

  6. mnist = input_data.read_data_sets("../MNIST_data/", one_hot=True)
  7. origin_test = mnist.test.images
  8. reshape_test = []
  9. for t in origin_test:
  10.     b = t.reshape(28,28)
  11.     reshape_test.append(b)
  12. for length in [100,500,1000,10000]:
  13.     with tf.Graph().as_default():
  14.         output_graph_def = tf.GraphDef()
  15.         output_graph_path = './mnist_frozen_graph.pb'

  17.         with open(output_graph_path, 'rb') as f:
  18.             output_graph_def.ParseFromString(f.read())
  19.             _ = tf.import_graph_def(output_graph_def, name="")
  20.      
  21.         with tf.Session() as sess:
  22.             sess.run(tf.global_variables_initializer())
  23.             input = sess.graph.get_tensor_by_name("x:0")
  24.             output = sess.graph.get_tensor_by_name("y_conv:0")
  25.             y_conv_2 = sess.run(output, feed_dict={input:reshape_test[0:length]})
  26.             y_2 = mnist.test.labels[0:length]
  27.             print("first image:",y_conv_2[0])
  28.             correct_prediction_2 = tf.equal(tf.argmax(y_conv_2, 1), tf.argmax(y_2, 1))
  29.             accuracy_2 = tf.reduce_mean(tf.cast(correct_prediction_2, "float"))
  30.             print('%d:'%length,"check accuracy %g" % sess.run(accuracy_2))
复制代码rknn_predict.py
  1. import numpy as np
  2. from PIL import Image
  3. from rknn.api import RKNN
  4. import cv2
  5. from tensorflow.examples.tutorials.mnist import input_data
  6. import tensorflow as tf

  8. mnist = input_data.read_data_sets("../MNIST_data/", one_hot=True)
  9. print(mnist.test.images[0].shape)
  10. # 解析模型的输出，获得概率最大的手势和对应的概率
  11. def get_predict(probability):
  12.     data = probability[0][0]
  13.     data = data.tolist()
  14.     max_prob = max(data)
  15.     return data.index(max_prob), max_prob
  16. # return data.index(max_prob), max_prob;
  17. def load_model(model_name):
  18.     # 创建RKNN对象
  19.     rknn = RKNN()
  20.     # 载入RKNN模型
  21.     print('-->loading model')
  22.     rknn.load_rknn(model_name)
  23.     print('loading model done')
  24.     # 初始化RKNN运行环境
  25.     print('--> Init runtime environment')
  26.     ret = rknn.init_runtime()
  27.     if ret != 0:
  28.        print('Init runtime environment failed')
  29.        exit(ret)
  30.     print('done')
  31.     return rknn
  32. def predict(rknn,length):
  33.     acc_count = 0
  34.     for i in range(length):
  35.         # im = mnist.test.images[i]
  36.         im = Image.open("../MNIST_data/mnist_test/test_%d.jpg"%i)   # 加载图片
  37.         im = im.resize((28,28),Image.ANTIALIAS)
  38.         im = np.asarray(im)
  39.         outputs = rknn.inference(inputs=[im])
  40.         pred, prob = get_predict(outputs)
  41.         if i ==0:
  42.             print(outputs)
  43.             print(prob)
  44.             print(pred)
  45.         if i ==100 or i ==500 or i ==1000 or i ==10000:
  46.             result = float(acc_count)/i
  47.             print('result%d:'%i,result)
  48.         if list(mnist.test.labels[i]).index(1) == pred:
  49.             acc_count += 1
  50.     result = float(acc_count)/length
  51.     print('result:',result)
  52.     # acc_count = 0
  53.     # length = len(mnist.test.images)
  54.     # for i in range(length):
  55.         # im = mnist.test.images[i]# 加载图片
  56.         # outputs = rknn.inference(inputs=[im])   # 运行推理，得到推理结果
  57.         # pred, prob = get_predict(outputs)     # 将推理结果转化为可视信息
  58.         # if i%100 == 0:
  59.             # print(prob)
  60.             # print(pred)
  61.             # print(acc_count)
  62.             # print(list(mnist.test.labels[i]).index(1))
  63.         # if list(mnist.test.labels[i]).index(1) == pred:
  64.             # acc_count += 1
  65.     # result = float(acc_count)/length
  66.     # print('result:',result)
  67. if __name__=="__main__":
  68.     #此处要改成相应的量化或者非量化rknn模型
  69.     model_name = './mnist.rknn'
  70.     length = 10000
  71.     rknn = load_model(model_name)
  72.     predict(rknn,length)

  74.     rknn.release()
复制代码得到最终的结果对比图表如下:


本帖子中包含更多资源

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

x
回复

使用道具 举报

peng

注册会员

积分
169
沙发
 楼主| 发表于 2019-9-20 08:43:10 | 显示全部楼层
本帖最后由 peng 于 2019-9-27 14:50 编辑

《量化训练》关键的两个函数为：tf.contrib.quantize.create_training_graph，tf.contrib.quantize.create_eval_graph(input_graph=g)，在上一节中代码中已经提供了，具体可看以及几个步骤：第一步，修改train.py,将create_training_graph()的is_quantify设为True.然后重新训练模型。
  1. if __name__ == "__main__":
  2. ckpt = './checkpoint_fake/mnist.ckpt'
  3. point_dir = './checkpoint_fake'
  4. pbtxt = 'mnist_fake.pbtxt'
  5. g1 = create_training_graph(True)
  6. train_network(g1,ckpt,point_dir,pbtxt)
复制代码第二步，修改freese.py,将frozen()的is_quantify设为True.重新固化成pb文件，
  1. if __name__ == "__main__":
  2. ckpt = './checkpoint_fake/mnist.ckpt'
  3. pbtxt = 'mnist_frozen_graph_fake.pb'
  4. frozen(True,ckpt,pbtxt)
复制代码
第三步，使用toco工具将pb模型转化为全量化的tflite模型(输入，输出和权值都变成int8类型)

  1. #!/bin/sh

  3. toco \
  4. --graph_def_file=mnist_frozen_graph_fake.pb \
  5. --output_file=mnist_fakequantize.tflite \
  6. --output_format=TFLITE \
  7. --inference_type=QUANTIZED_UINT8 \
  8. --input_shapes=1,28,28 \
  9. --input_arrays=x \
  10. --output_arrays=y_conv \
  11. --mean_values=0 \
  12. --std_dev_values=256 \
  13. --change_concat_input_ranges=false --allow_custom_ops
复制代码

第四步，将tflite模型转为rknn，修改rknn_transfer.py:

  1. def tflite_transfer():
  2. rknn = RKNN()
  3. print('--> Loading model')
  4. ret = rknn.load_tflite(model = './mnist_fakequantize.tflite')
  5. print('done')
  6. print('--> Building model')
  7. rknn.build(do_quantization=False)
  8. print('done')
  9. # 导出保存rknn模型文件
  10. rknn.export_rknn('./mnist_quantization_fake.rknn')
  11. # Release RKNN Context
  12. rknn.release()
复制代码
  1. if __name__ == '__main__':
  2. tflite_transfer()
复制代码第五步，分别运行tflite_predict.py,和rknn_predict.py，对比tflite和rknn推理的结果和准确度：
tflite_predict.py
  1. mport tensorflow as tf
  2. from tensorflow.examples.tutorials.mnist import input_data
  3. from PIL import Image

  5. mnist = input_data.read_data_sets("../../MNIST_data/", one_hot=True)
  6. length = 100
  7. # 加载模型并分配张量
  8. interpreter = tf.lite.Interpreter(model_path="mnist_fakequantize.tflite")
  9. interpreter.allocate_tensors()
  10. # 获取输入输出张量
  11. input_details = interpreter.get_input_details()
  12. print(input_details)
  13. output_details = interpreter.get_output_details()
  14. print(output_details)



  18. acc_count = 0
  19. for i in range(length):
  20. #im = mnist.test.images[i]
  21. im = Image.open("../../MNIST_data/mnist_test/test_%d.jpg"%i) # 加载图片
  22. im = im.resize((28,28),Image.ANTIALIAS)
  23. im = np.asarray(im)
  24. # print(im.dtype)
  25. input_shape = input_details[0]['shape']
  26. input_data = im.reshape(1,28,28)
  27. # print(input_data.dtype)
  28. interpreter.set_tensor(input_details[0]['index'], input_data)

  30. interpreter.invoke()
  31. output_data = interpreter.get_tensor(output_details[0]['index'])
  32. # print(output_data)
  33. # print(output_data.argmax(axis=1)[0])
  34. if i <3:
  35. print(output_data)

  37. if list(mnist.test.labels[i]).index(1) == output_data.argmax(axis=1)[0]:
  38. acc_count += 1
  39. result = float(acc_count)/length
  40. print('result:',result)
复制代码rknn_predict.py
  1. import numpy as np
  2. from PIL import Image
  3. from rknn.api import RKNN
  4. import cv2
  5. from tensorflow.examples.tutorials.mnist import input_data
  6. import tensorflow as tf

  8. mnist = input_data.read_data_sets("../MNIST_data/", one_hot=True)
  9. print(mnist.test.images[0].shape)
  10. # 解析模型的输出，获得概率最大的手势和对应的概率
  11. def get_predict(probability):
  12. data = probability[0][0]
  13. data = data.tolist()
  14. max_prob = max(data)
  15. return data.index(max_prob), max_prob
  16. # return data.index(max_prob), max_prob;
  17. def load_model(model_name):
  18. # 创建RKNN对象
  19. rknn = RKNN()
  20. # 载入RKNN模型
  21. print('-->loading model')
  22. rknn.load_rknn(model_name)
  23. print('loading model done')
  24. # 初始化RKNN运行环境
  25. print('--> Init runtime environment')
  26. ret = rknn.init_runtime()
  27. if ret != 0:
  28. print('Init runtime environment failed')
  29. exit(ret)
  30. print('done')
  31. return rknn
  32. def predict(rknn,length):
  33. acc_count = 0
  34. for i in range(length):
  35. im = mnist.test.images[i]
  36. # im = Image.open("../MNIST_data/mnist_test/test_%d.jpg"%i) # 加载图片
  37. # im = im.resize((28,28),Image.ANTIALIAS)
  38. # im = np.asarray(im)
  39. im = im.reshape(1,28,28)
  40. outputs = rknn.inference(inputs=[im])
  41. pred, prob = get_predict(outputs)
  42. if i ==0:
  43. print(outputs)
  44. print(prob)
  45. print(pred)
  46. if i ==100 or i ==500 or i ==1000 or i ==10000:
  47. result = float(acc_count)/i
  48. print('result%d:'%i,result)
  49. if list(mnist.test.labels[i]).index(1) == pred:
  50. acc_count += 1
  51. result = float(acc_count)/length
  52. print('result:',result)
  53. if __name__=="__main__":
  54. #此处要改成相应的量化或者非量化rknn模型
  55. model_name = './mnist_quantization_fake.rknn'
  56. length = 10000
  57. rknn = load_model(model_name)
  58. predict(rknn,length)

  60. rknn.release()
复制代码最终运行结果如下表所示：

本帖子中包含更多资源

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

x
回复

使用道具 举报

返回列表
高级模式
B Color Image Link Quote Code Smilies
您需要登录后才可以回帖 登录 | 立即注册

本版积分规则

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

瑞芯微电子股份有限公司所有       闽ICP备19006074号-1      闽公网安备 35010202001061号

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