## 获得随机数
def generateRandomInput(batch_size=16):
xs= np.random.randn(batch_size, 784)
ys=np.zeros((batch_size,10))
for i in range(batch_size):
ys[i][np.random.randint(0, 10)]=1
return xs,ys
## 构造训练数据
def getFixedInput(reuse=None):
if reuse:
with open("test.txt", "r", encoding="utf-8") as f:
xsInputs=[]
ysInputs=[]
for line in f.readlines():
d=json.loads(line)
xsInputs.append([float(i) for i in d["xs"].split(",")])
ysInputs.append([float(i) for i in d["ys"].split(",")])
return np.array(xsInputs),np.array(ysInputs)
xs,ys=generateRandomInput(10000)
with open("test.txt","w",encoding="utf-8") as f:
for (x,y) in zip(xs,ys):
d=dict()
d["xs"]=",".join([str(i) for i in list(x)])
d["ys"]=",".join([str(i) for i in list(y)])
json.dump(d, f)
f.write('n')
return xs,ys
## 构造训练数据
def getFixedInput(reuse=None):
if reuse:
with open("test.txt", "r", encoding="utf-8") as f:
xsInputs=[]
ysInputs=[]
for line in f.readlines():
d=json.loads(line)
xsInputs.append([float(i) for i in d["xs"].split(",")])
ysInputs.append([float(i) for i in d["ys"].split(",")])
return np.array(xsInputs),np.array(ysInputs)
xs,ys=generateRandomInput(10000)
with open("test.txt","w",encoding="utf-8") as f:
for (x,y) in zip(xs,ys):
d=dict()
d["xs"]=",".join([str(i) for i in list(x)])
d["ys"]=",".join([str(i) for i in list(y)])
json.dump(d, f)
f.write('n')
return xs,ys
## 模型生成和onnx模型保存
def train_model(model_dir):
tf.reset_default_graph()
config = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True))
config.gpu_options.per_process_gpu_memory_fraction = 0.99
with tf.Session(config=config) as sess:
## 定义x,y,keep_prob占位符
xs = tf.placeholder(tf.float32, [None, 784],name="xs")
ys = tf.placeholder(tf.float32, [None, 10],name="ys")
keep_prob = tf.placeholder(tf.float32,name="keep_prob")
# keep_prob = tf.placeholder_with_default(1.0, shape=(), name="keep_prob") # dropout (keep probability
x_image = tf.reshape(xs, [-1, 28, 28, 1]) # [n_samples,28,28,1] -1代表输入暂时不定可以多张输入
with tf.variable_scope("conv1_layer"):
### conv1 layer ###
W_conv1 = weight_variable([5, 5, 1, 32],) ##patch 5x5 ,in size 1,out size 32
b_conv1 = biases_variable([32])
h_conv1 = tf.nn.relu(conv2d(x_image, W_conv1) + b_conv1) # out size 28x28x32
h_pool1 = max_pool_2x2(h_conv1) # out size 14x14x32
with tf.variable_scope("conv2_layer"):
###conv2 layer ###
W_conv2 = weight_variable([5, 5, 32, 64]) # patch 5x5 ,in size 32 ,out size 64
b_conv2 = biases_variable([64])
h_conv2 = tf.nn.relu(conv2d(h_pool1, W_conv2) + b_conv2) # out size 14x14x32
h_pool2 = max_pool_2x2(h_conv2) # output size 7x7x64
with tf.variable_scope("fc1_layer"):
### fc1 layer###
W_fc1 = weight_variable([7 * 7 * 64, 1024])
b_fc1 = biases_variable([1024])
## [n_samples ,7 ,7 ,64] ->> [n_samples ,7x7x64]
h_pool2_flat = tf.reshape(h_pool2, [-1, 7 * 7 * 64])
h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, W_fc1) + b_fc1)
h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)
# h_fc1_drop=h_fc1
with tf.variable_scope("fc2_layer"):
### fc2 layer###
W_fc2 = weight_variable([1024, 10])
b_fc2 = biases_variable([10])
prediction = tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2,name="predict")
tf.identity(prediction, name="output")
loss = tf.reduce_mean(-tf.reduce_sum(ys * tf.log(prediction), reduction_indices=[1]))
train_step = tf.train.AdamOptimizer(1e-4).minimize(loss)
init = tf.initialize_all_variables()
sess.run(init)
# for i in range(1000):
batch_xs, batch_ys = generateRandomInput(16)
# print(type(batch_xs), batch_ys.shape)
sess.run(train_step, feed_dict={xs: batch_xs, ys: batch_ys, keep_prob: 0.5})
# sess.run(train_step, feed_dict={xs: batch_xs, ys: batch_ys})
node_names_list = [tensor.name for tensor in tf.get_default_graph().as_graph_def().node]
# for node_name in node_names_list:
# print(node_name)
saver = tf.train.Saver() # 实例化saver对象，用于保存模型
saver.save(sess, '{}/mymodel'.format(model_dir))
save_pb_model(sess, '{}/'.format(model_dir), ['output'])
# print('my model saved in file:', my_model)
frozen_graph = freeze_session(sess, input_names=["xs:0","keep_prob:0"],output_names=["output:0"]) ##模型冻结，将参数固化
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(frozen_graph, name='')
with tf.Session(graph=tf_graph) as sess:
onnx_graph = tf2onnx.tfonnx.process_tf_graph(sess.graph, input_names=["xs:0","keep_prob:0"], output_names=["output:0"])
# onnx_graph = tf2onnx.tfonnx.process_tf_graph(sess.graph, input_names=["xs:0"],output_names=["output:0"])
model_proto = onnx_graph.make_model("model")
with open("{}/model.onnx".format(model_dir), "wb") as f: ##onnx模型保存
f.write(model_proto.SerializeToString())
print("done")
##onnx 模型加载
def myCheckONNX( xs, ys ,model_dir):
print("myCheckONNX")
onnx_model = onnx.load("{}/model.onnx".format(model_dir))
onnx.checker.check_model(onnx_model)
import onnxruntime
ort_session = onnxruntime.InferenceSession("{}/model.onnx".format(model_dir),providers=['CUDAExecutionProvider'])
def to_numpy(tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
# compute ONNX Runtime output prediction
# ort_inputs = {ort_session.get_inputs()[0].name: input1,ort_session.get_inputs()[1].name:d[0]}
print(ort_session.get_inputs()[1].name)
a = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9], dtype='float64')
c = a.astype(np.float32)
ort_inputs = {ort_session.get_inputs()[0].name: xs.astype(np.float32),ort_session.get_inputs()[1].name: np.array((1.0)).astype(np.float32)}
ort_outs = ort_session.run(["output:0"], ort_inputs)
print(ort_outs)
return ort_outs[0]
if __name__ == '__main__':
model_dir="models_keepProb"
train_model(model_dir)
xs, ys = getFixedInput()
res_onnx=myCheckONNX(xs, ys,model_dir)
