|
|
本帖最后由 jiajia1990 于 2019-8-27 13:47 编辑
在电脑上训练好PB文件后，在开发板上转换为rknn文件，并推理，但不正确。不知道哪里转换的不对？
以下是tensorflow训练代码，pb文件生成后，测试正常。
import tensorflow as tf
import tensorflow.examples.tutorials.mnist.input_data as input_data
from tensorflow.python.framework import graph_util
import time
from tensorflow.python.platform import gfile
def weight(shape):
return tf.Variable(tf.truncated_normal(shape, stddev=0.1), name='W')
def bias(shape):
return tf.Variable(tf.constant(0.1, shape=shape, name='b'))
def conv2d(x, W):
return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')
def max_pool_2x2(x):
return tf.nn.max_pool(x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
mnist = input_data.read_data_sets('MNIST_data/', one_hot=True)
def train():
with tf.name_scope('Input_Layer'):
x = tf.placeholder("float", shape=[None, 784], name='x')
print(x.name)
x_image = tf.reshape(x, [-1, 28, 28, 1])
with tf.name_scope('C1_Conv'):
W1 = weight([5, 5, 1, 16])
b1 = bias([16])
Conv1 = conv2d(x_image, W1) + b1
C1_Conv = tf.nn.relu(Conv1)
with tf.name_scope('C1_Pool'):
C1_Pool = max_pool_2x2(C1_Conv)
with tf.name_scope('C2_Conv'):
W2 = weight([5, 5, 16, 36])
b2 = bias([36])
Conv2 = conv2d(C1_Pool, W2) + b2
C2_Conv = tf.nn.relu(Conv2)
with tf.name_scope('C2_Pool'):
C2_Pool = max_pool_2x2(C2_Conv)
with tf.name_scope('D_Flat'):
D_Flat = tf.reshape(C2_Pool, [-1, 1764])
with tf.name_scope('D_Hidden_Layer'):
W3 = weight([1764, 128])
b3 = bias([128])
D_Hidden = tf.nn.relu(tf.matmul(D_Flat, W3) + b3)
D_Hidden_Dropout = tf.nn.dropout(D_Hidden, keep_prob=0.8)
with tf.name_scope('Output_Layer'):
W4 = weight([128, 10])
b4 = bias([10])
y_predict = tf.nn.softmax(tf.matmul(D_Hidden_Dropout, W4) + b4, name='Output_Layer')
print("y_predict is :", y_predict)
with tf.name_scope('optimizer'):
y_label = tf.placeholder("float", shape=[None, 10], name="y_label")
loss_function = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=y_predict, labels=y_label))
optimizer = tf.train.AdamOptimizer(learning_rate=0.0001).minimize(loss_function)
with tf.name_scope('evaluate_model'):
correct_prediction = tf.equal(tf.argmax(y_predict, 1), tf.argmax(y_label, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, 'float'))
trainEpochs = 15
batchSize = 100
totalBatchs = 550
loss_list = []
epoch_list = []
accuracy_list = []
startTime = time.time()
sess = tf.Session()
sess.run(tf.global_variables_initializer())
for epoch in range(trainEpochs):
for i in range(totalBatchs):
batch_x, batch_y = mnist.train.next_batch(batchSize)
sess.run(optimizer, feed_dict={x: batch_x, y_label: batch_y})
loss, acc = sess.run([loss_function, accuracy], feed_dict={x: mnist.validation.images, y_label: mnist.validation.labels})
epoch_list.append(epoch)
loss_list.append(loss)
accuracy_list.append(acc)
print('Train Epoch:', '%02d' % (epoch+1), 'Loss=', '{:.9f}'.format(loss), "Accuracy=", acc)
endTime = time.time()
print("time is %d" % (endTime - startTime))
graph_def = tf.get_default_graph().as_graph_def()
output_graph_def = graph_util.convert_variables_to_constants(sess, graph_def, ['Output_Layer/Output_Layer'])
print(output_graph_def, file=open("log/output_graph_def.txt", 'w'))
with tf.gfile.GFile("./model.pb", "wb") as f:
f.write(output_graph_def.SerializeToString())
print("Accuracy:", sess.run(accuracy, feed_dict={x: mnist.test.images, y_label: mnist.test.labels}))
prediction_result = sess.run(tf.argmax(y_predict, 1), feed_dict={x: mnist.test.images})
print(prediction_result[:5])
print(mnist.test.labels[:5])
# plot_images_labels_prediction(mnist.test.images, mnist.test.labels, prediction_result, 0)
sess.close()
以下是pb转rknn的代码。
from tensorflow.examples.tutorials.mnist import input_data
from rknn.api import RKNN
INPUT_SIZE = 784
if __name__ == '__main__':
# 创建RKNN执行对象
rknn = RKNN()
rknn.config(channel_mean_value='0 0 0 255')
print('--> Loading model')
rknn.load_tensorflow(tf_pb='./model.pb',
inputs=['Input_Layer/x'],
outputs=['Output_Layer/Output_Layer'],
input_size_list=[[INPUT_SIZE]])
print('done')
# 创建解析pb模型
print('--> Building model')
rknn.build(do_quantization=False)
print('done')
# 导出保存rknn模型文件
rknn.export_rknn('./model.rknn')
# Release RKNN Context
rknn.release()
以下是rknn推理代码：
import numpy as np
from PIL import Image
from tensorflow.examples.tutorials.mnist import input_data
from rknn.api import RKNN
def get_predict(probability):
data = probability[0][0]
data = data.tolist()
max_prob = max(data)
return data.index(max_prob), max_prob
def load_model():
# 创建RKNN对象
rknn = RKNN()
# 载入RKNN模型
print('-->loading model')
rknn.load_rknn('./model.rknn')
print('loading model done')
# 初始化RKNN运行环境
print('--> Init runtime environment')
# ret = rknn.init_runtime(target='rk3399pro')
ret = rknn.init_runtime()
if ret != 0:
print('Init runtime environment failed')
exit(ret)
print('done')
return rknn
def predict(rknn):
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True) # 载入数据集
images, labels = mnist.test.next_batch(1)
print(labels)
print(images.shape)
outputs = rknn.inference(inputs=[images]) # 运行推理，得到推理结果
print(outputs)
pred, prob = get_predict(outputs) # 将推理结果转化为可视信息
print(prob)
print(pred)
if __name__ == "__main__":
rknn = load_model()
predict(rknn)
rknn.release()
推理出来的每次都不对，实际为[[0. 0. 0. 0. 0. 1. 0. 0. 0. 0.]]，推理出来是[array([[0. , 0. , 0.5, 0. , 0. , 0. , 0. , 0.5, 0. , 0. ]], dtype=float32)]样子的，pb文件使用tensorflow在电脑上测试是对的，但转成rknn就不对了，请问是哪里出的问题吗？之前看一个人发的帖子也遇到这个问题了，但他没有说怎么解决的。
|
|
收藏
闽公网安备 35010202001061号