TensorFlow/TensorFlow 保存 pb 模型

TensorFlow/TensorFlow 保存 pb 模型

我们在保存 TensorFlow 的模型时,一般都是把模型保存为 ckpt 文件。

但是如果把模型迁移到手机上,那么就要把模型保存为 pb 文件。模型保存为 pb 文件时候,模型的变量都会变成固定的,从而会大大减小模型的大小,适合在手机端运行。

这里来讲下如何把模型保存为 pb 文件,以及如何从 pb 文件中加载模型。

pb 的保存

下面是一个例子。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import tensorflow as tf
from tensorflow.python.framework import graph_util

v1 = tf.Variable(tf.constant(1.0, shape=[1]), name="v1")
v2 = tf.Variable(tf.constant(2.0, shape=[1]), name="v2")
result = v1 + v2

init_op = tf.global_variables_initializer()
with tf.Session() as sess:
sess.run(init_op) # 初始化所有变量
# 导出当前计算图的GraphDef部分
graph_def = tf.get_default_graph().as_graph_def()
# 保存对应计算节点的名称。
output_graph_def = graph_util.convert_variables_to_constants(sess, graph_def, ['add'])
# 将导出的模型存入文件中
with tf.gfile.GFile("model.pb", "wb") as f:
f.write(output_graph_def.SerializeToString())

下面的代码中,计算了 \(result = v1 + v2\)

但我只想保存 result 变量,关键的代码是

1
output_graph_def = graph_util.convert_variables_to_constants(sess, graph_def, ['add'])

其中 ['add'] 是一个 list,里面存放的是要保存的计算节点的名称,因为我们使用的加法,所以计算节点的名称就是 add

同理,如果你想保存乘法的计算节点,名称是 multiply

你还可以给每个计算节点起名称:

例如上面的

1
result = v1 + v2

默认的名称是 add,可以使用 TensorFlow 提供的 tf.add() 来替代,并自定义名称:

1
result = tf.add(v1,v2,name='result')

那么在保存的时候,add 就改为 result

1
output_graph_def = graph_util.convert_variables_to_constants(sess, graph_def, ['result'])

如果你想查看所有计算节点的名称,可以使用下面的语句:

1
var_list = tf.global_variables()

var_list 就保存了一部分张量的名称

或者直接查看某个变量的 name 属性,例如查看 result 的张量名称:

1
2
3
result = tf.multiply(v1,v2,name='result')

print(result.name)

输出是 result:0。(注意,这里是 result:0,而不是 result

这里需要注意,张量的名称计算节点的名称不是一个概念。它们是从属关系。具体来说,一个计算节点(也被称为一个Operation)可以包括多个张量。具体到名称上,张量的名称 = 计算节点的名称 + ":"+编号

因此,上面打印了 result.name 的值是 result:0,我们要取出前面的 result,作为计算节点的名称。

1
output_graph_def = graph_util.convert_variables_to_constants(sess, graph_def, ['result'])

pb 的加载

代码如下:

1
2
3
4
5
6
7
8
9
10
11
from tensorflow.python.platform import gfile

with tf.Session() as sess:
model_filename = "model.pb"
# 读取保存的模型文件,并将其解析成对应的GraphDef Protocol Buffer
with gfile.FastGFile(model_filename, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
# 将graph_def中保存的图加载到当前图中,注意读取的是张量的名称
result = tf.import_graph_def(graph_def, return_elements=["result:0"])
print(sess.run(result))

注意这里读取的是张量的名称 "result:0"

下面使用一个实际的案例来讲解保存和加载 pb 模型。

保存 pb 模型

在这个例子中,我们定义了一个卷积神经网络。

一般来说,我们只需要保存输入的 placeholder,以及最终输出的结果、loss 等数据。

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
num_classes=20
input_image = tf.placeholder(tf.float32, shape=(None, 128, 128, 2), name='input')
kernel_size =5
layer1_conv = tf.layers.conv2d(
input_image,
64,
kernel_size=(kernel_size, kernel_size),
strides=(2, 2),
padding="SAME",
activation=tf.nn.relu,
name="layer1_conv",
)
layer2_conv = tf.layers.conv2d(
layer1_conv,
128,
kernel_size=(kernel_size, kernel_size),
strides=(2, 2),
padding="SAME",
activation=tf.nn.relu,
name="layer2_conv",
)
layer3_conv = tf.layers.conv2d(
layer2_conv,
256,
kernel_size=(kernel_size, kernel_size),
strides=(2, 2),
padding="SAME",
activation=tf.nn.relu,
name="layer3_conv",
)
layer4_conv = tf.layers.conv2d(
layer3_conv,
512,
kernel_size=(kernel_size, kernel_size),
strides=(2, 2),
padding="SAME",
activation=tf.nn.relu,
name="layer4_conv",
)

layer5_conv1 = tf.layers.conv2d(
layer4_conv,
512,
kernel_size=(kernel_size, kernel_size),
strides=(1, 1),
padding="SAME",
activation=tf.nn.relu,
name="layer5_conv1",
)
layer5_conv2 = tf.layers.conv2d(
layer5_conv1,
512,
kernel_size=(kernel_size, kernel_size),
strides=(1, 1),
padding="SAME",
activation=tf.nn.relu,
name="layer5_conv2",
)
layer5_res = tf.add(layer5_conv2, layer4_conv, name="layer5_res")
layer5_relu2 = tf.nn.relu(layer5_res, name="layer5_relu2")

layer6_conv1 = tf.layers.conv2d(
layer5_relu2,
512,
kernel_size=(kernel_size, kernel_size),
strides=(1, 1),
padding="SAME",
activation=tf.nn.relu,
name="layer6_conv1",
)
layer6_conv2 = tf.layers.conv2d(
layer6_conv1,
512,
kernel_size=(kernel_size, kernel_size),
strides=(1, 1),
padding="SAME",
activation=tf.nn.relu,
name="layer6_conv2",
)
layer6_res = tf.add(layer6_conv2, layer5_relu2, name="layer6_res")
layer6_relu2 = tf.nn.relu(layer6_res, name="layer6_relu2")

layer7_input = tf.reshape(layer6_relu2, [-1, 8 * 8 * 512])

layer7_dense = tf.layers.dense(
layer7_input, 1024, activation=tf.nn.relu, name="layer7_dense"
)
layer8_dense = tf.layers.dense(
layer7_dense, 512, activation=tf.nn.relu, name="layer8_dense"
)
layer9_dense = tf.layers.dense(layer8_dense, num_classes, name="layer9_dense")


input_lable = tf.placeholder(tf.float32, shape=(None, num_classes), name='label')

probability = tf.nn.softmax(logits=layer9_dense, name="probability")

total_loss = tf.losses.softmax_cross_entropy(
onehot_labels=input_lable, logits=layer9_dense
)

optimizer = tf.train.AdamOptimizer(learning_rate=0.00003).minimize(total_loss)

init = tf.initialize_all_variables()
sess = tf.Session()
sess.run(init)

train_images = np.ones([1, 128, 128, 2], dtype = int)
train_labels = np.zeros([1, num_classes], dtype = int)
train_labels[0][1]=1

# 训练
_, loss_result, prob_result = sess.run(
[optimizer, total_loss, probability],
{input_image: train_images, input_lable: train_labels},
)

print(loss_result)
# 2.9825938


# 导出当前计算图的GraphDef部分
graph_def = tf.get_default_graph().as_graph_def()
# 查看所有节点的名称
var_list = tf.global_variables()
# 保存指定的节点,并将节点值保存为常数
output_graph_def = graph_util.convert_variables_to_constants(sess, graph_def, ['softmax_cross_entropy_loss/value','input', 'label', "probability"])
# 将计算图写入到模型文件中
model_f = tf.gfile.GFile("cnn_model.pb", "wb")
model_f.write(output_graph_def.SerializeToString())

在上面的代码中,我们保存的计算节点包括:['softmax_cross_entropy_loss/value','input', 'label', "probability"],分别对应 loss、输入、标签以及最终输出的概率,而不用关心中间的那些网络层。

加载 pb 模型

下面是加载 pb 模型的代码。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
sess = tf.Session()
# 将保存的模型文件解析为GraphDef
model_f = gfile.FastGFile("cnn_model.pb", 'rb')
graph_def = tf.GraphDef()
graph_def.ParseFromString(model_f.read())
tf.import_graph_def(graph_def, name="")
input_image = sess.graph.get_tensor_by_name("input:0")
input_lable = sess.graph.get_tensor_by_name("label:0")
total_loss = sess.graph.get_tensor_by_name("softmax_cross_entropy_loss/value:0")
probability = sess.graph.get_tensor_by_name("probability:0")
train_images = np.ones([1, 128, 128, 2], dtype = int)
train_labels = np.zeros([1, 20], dtype = int)
train_labels[0][1]=1
loss_result, prob_result = sess.run(
[total_loss, probability],
{input_image: train_images, input_lable: train_labels},
)

print(loss_result)

其中关键的代码是:

1
2
3
4
input_image = sess.graph.get_tensor_by_name("input:0")
input_lable = sess.graph.get_tensor_by_name("label:0")
total_loss = sess.graph.get_tensor_by_name("softmax_cross_entropy_loss/value:0")
probability = sess.graph.get_tensor_by_name("probability:0")

分别把 loss、输入、标签以及最终输出的概率取出来,注意这里使用的是张量名称。然后就可以进行计算和推理了。

你学会了吗?

如果你还有不明白的地方,欢迎给我留言。

如果你觉得这篇文章对你有帮助,不妨给我点个赞,鼓励我写出更多好文章。

参考资料

评论