-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathtransform_cpu.py
134 lines (117 loc) · 5.69 KB
/
transform_cpu.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
# coding:utf-8
import os
import tensorflow as tf
from tensorflow.python.platform import gfile
import argparse
import numpy as np
import cv2 as cv
import time
from power_diff_numpy import *
os.putenv('MLU_VISIBLE_DEVICES','')
def parse_arg():
parser = argparse.ArgumentParser()
parser.add_argument('image')
parser.add_argument('ori_pb')
parser.add_argument('ori_power_diff_pb')
parser.add_argument('numpy_pb')
args = parser.parse_args()
return args
def run_ori_pb(ori_pb, image):
config = tf.ConfigProto(allow_soft_placement=True,
inter_op_parallelism_threads=1,
intra_op_parallelism_threads=1)
model_name = os.path.basename(ori_pb).split(".")[0]
image_name = os.path.basename(image).split(".")[0]
g = tf.Graph()
with g.as_default():
with tf.gfile.FastGFile(ori_pb,'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
tf.import_graph_def(graph_def, name='')
img = cv.imread(image)
X = cv.resize(img, (256, 256))
with tf.Session(config=config) as sess:
sess.graph.as_default()
sess.run(tf.global_variables_initializer())
input_tensor = sess.graph.get_tensor_by_name('X_content:0')
output_tensor = sess.graph.get_tensor_by_name('add_37:0')
start_time = time.time()
ret =sess.run(output_tensor, feed_dict={input_tensor:[X]})
end_time = time.time()
print("C++ inference(CPU) origin pb time is: ",end_time-start_time)
img1 = tf.reshape(ret,[256,256,3])
img_numpy = img1.eval(session=sess)
cv.imwrite(image_name + '_' + model_name + '_cpu.jpg',img_numpy)
def run_ori_power_diff_pb(ori_power_diff_pb, image):
config = tf.ConfigProto(allow_soft_placement=True,
inter_op_parallelism_threads=1,
intra_op_parallelism_threads=1)
model_name = os.path.basename(ori_power_diff_pb).split(".")[0]
image_name = os.path.basename(image).split(".")[0]
g = tf.Graph()
with g.as_default():
with tf.gfile.FastGFile(ori_power_diff_pb,'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
tf.import_graph_def(graph_def, name='')
img = cv.imread(image)
X = cv.resize(img, (256, 256))
with tf.Session(config=config) as sess:
# 完成PowerDifference Pb模型的推理
sess.graph.as_default()
sess.run(tf.global_variables_initializer())
input_tensor = sess.graph.get_tensor_by_name('X_content:0')
input_tensor_pow = sess.graph.get_tensor_by_name('moments_15/PowerDifference_z:0')
output_tensor = sess.graph.get_tensor_by_name('add_37:0')
input_pow = np.array(2, dtype=float)
start_time = time.time()
ret = sess.run(output_tensor, feed_dict={input_tensor:[X], input_tensor_pow: input_pow})
end_time = time.time()
print("C++ inference(CPU) time is: ",end_time-start_time)
img1 = tf.reshape(ret,[256,256,3])
img_numpy = img1.eval(session=sess)
cv.imwrite(image_name + '_' + model_name + '_cpu.jpg',img_numpy)
def run_numpy_pb(numpy_pb, image):
config = tf.ConfigProto(allow_soft_placement=True,
inter_op_parallelism_threads=1,
intra_op_parallelism_threads=1)
model_name = os.path.basename(numpy_pb).split(".")[0]
image_name = os.path.basename(image).split(".")[0]
g = tf.Graph()
with g.as_default():
with tf.gfile.FastGFile(numpy_pb,'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
tf.import_graph_def(graph_def, name='')
img = cv.imread(image)
X = cv.resize(img, (256, 256))
with tf.Session(config=config) as sess:
# 完成Numpy版本 Pb模型的推理
sess.graph.as_default()
sess.run(tf.global_variables_initializer())
# 根据输入名称获得输入的tensor
input_tensor = sess.graph.get_tensor_by_name('X_content:0')
# 获取两个输出节点， 作为numpy算子的输入
out_tmp_tensor_1 = sess.graph.get_tensor_by_name('Conv2D_13:0')
out_tmp_tensor_2 = sess.graph.get_tensor_by_name('moments_15/StopGradient:0')
start_time = time.time()
# 执行第一次sessionrun， 得到numpy算子的两个输入值,注意此时两个输入的shape不同
input_x, input_y = sess.run([out_tmp_tensor_1, out_tmp_tensor_2], feed_dict={input_tensor:[X]})
input_pow = 2
output = power_diff_numpy(input_x, input_y, input_pow).reshape(1, 256, 256, 3)
# 根据添加的输入节点名称获得输入tensor
input_tensor_new = sess.graph.get_tensor_by_name('moments_15/PowerDifference:0')
# 完整推理最终输出的tensor
output_tensor = sess.graph.get_tensor_by_name('add_37:0')
# 执行第二次session run，输入图片数据以及上一步骤numpy计算的数据
ret = sess.run(output_tensor, feed_dict={input_tensor:[X], input_tensor_new: output})
end_time = time.time()
print("Numpy inference(CPU) time is: ",end_time-start_time)
img1 = tf.reshape(ret,[256,256,3])
img_numpy = img1.eval(session=sess)
cv.imwrite(image_name + '_' + model_name + '_cpu.jpg',img_numpy)
if __name__ == '__main__':
args = parse_arg()
run_ori_pb(args.ori_pb, args.image)
run_ori_power_diff_pb(args.ori_power_diff_pb, args.image)
run_numpy_pb(args.numpy_pb, args.image)