TensorFlow代码结构优化tips

TensorFlow代码结构优化tips

---

一、变量管理

在模型的前向传播的过程中，将前向传播定义为：

def inference(input_tensor, avg_class, weights1, biases1, weights2, biases2)

从上面定义可以看出，这个函数的参数包括了神经网络中的所有参数。然而，当神经网络的结构更加复杂、参数更多的时间，就需要一个更好的方式来传递和管理参数了。

TensorFlow提供了通过变量名称来创建或者获取一个变量的机制，通过这个机制，在不同的函数中可以直接通过变量的名字来使用变量，儿不需要将变量通过参数的形式到处传递。

变量管理机制主要是通过tf.get_variable和tf.variable_scope函数来实现的。

tf.get_variable 获取变量

在通过tf.get_variable 获取一个已经创建的变量时，需要通过tf.variable_scope函数来生成一个上下文管理器，并明确指定在这个上下文管理器中，tf.get_variable将直接获取已经生成的变量。

with tf.variable_scope("foo"):
    v = tf.get_variable("v", [1], initializer=tf.constant_initializer(1.0))

# 报错，因为在命名空间“foo”中已经存在了v变量，如果没有指定reuse=True则会报错
#with tf.variable_scope("foo"):
   # v = tf.get_variable("v", [1])

with tf.variable_scope("foo", reuse=True):
    v1 = tf.get_variable("v", [1])
print v == v1

# 报错，因为在命名空间"bar"中，没有创建过v变量，直接使用reuse=True，会报错。
#with tf.variable_scope("bar", reuse=True):
   # v = tf.get_variable("v", [1])

tf.variable_scope 上下文管理

当tf.variable_scope函数使用参数resue=True生成上下文管理器时，其内部所有的 tf.get_variable函数都会直接获取已经创建的变量，如果变量不存在，则会报错。

tf.variable_scopewith tf.variable_scope("root"):
    print tf.get_variable_scope().reuse

    with tf.variable_scope("foo", reuse=True):
        print tf.get_variable_scope().reuse

        with tf.variable_scope("bar"):
            print tf.get_variable_scope().reuse

    print tf.get_variable_scope().reuse

output:
False
True
True
False

通过tf.variable_scope来管理变量

tf.variable_scope函数除了可以控制tf.get_variable执行的功能之外，这个函数也提供了一个管理变量命名空间的方式。

v1 = tf.get_variable("v", [1])
print v1.name

with tf.variable_scope("foo",reuse=True):
    v2 = tf.get_variable("v", [1])
print v2.name

with tf.variable_scope("foo"):
    with tf.variable_scope("bar"):
        v3 = tf.get_variable("v", [1])
        print v3.name

v4 = tf.get_variable("v1", [1])
print v4.name

output:
v:0
foo/v:0
foo/bar/v:0
v1:0

我们还可以通过变量的名称来直接获取变量

with tf.variable_scope("",reuse=True):
    v5 = tf.get_variable("foo/bar/v", [1])
    print v5 == v3
    v6 = tf.get_variable("v1", [1])     
    print v6 == v4

output:
True
True

前向传播结果函数的改进

通过tf.variable_scope和tf.get_variable函数进行前向传播的改进：

def inference(input_tensor, reuse = False):

    # 定义第一层神经网络的变量和前向传播的过程
    with tf.variable_scope('layer1', reuse = reuse):
        weights = tf.get_variable("weights",[INPUT_NODE, LAYER1_NODE], initializer = tf.truncated_normal_initializer(stddev = 0.1))
        biases = tf.get_variable("biases",[LAYER1_NODE], initializer = tf.constant_initializer(0.0))
        layer1 = tf.nn.relu(tf.matmul(input_tensor, weights) + biases)

        # 定义第二层神经网络的变量和前向传播的过程
        with tf.variable_scope('layer2', reuse = reuse):
        weights = tf.get_variable("weights",[LAYER1_NODE, OUTPUT_NODE], initializer = tf.truncated_normal_initializer(stddev = 0.1))
        biases = tf.get_variable("biases",[OUTPUT_NODE], initializer = tf.constant_initializer(0.0))
        layer2 = tf.matmul(input_tensor, weights) + biases

    return layer2

x = tf.placeholder(tf.float32, [None, INPUT_NODE], name = 'x-input')
y = inference(x)

二、TensorFlow模型持久化

TensorFlow的持久化方法，可以通过程序来保存一个训练好的模型，并从持久化之后的模型文件中还原被保存的模型。

持久化代码实现

1. 保存和还原网络模型—tf.train.Saver

import tensorflow as tf

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()
saver = tf.train.Saver()

# 利用Saver类的save函数保存模型的结构和变量的取值，以及模型文件列表
with tf.Session() as sess:
    sess.run(init_op)
    saver.save(sess, "Saved_model/model.ckpt")

# 重新加载保存了两个变量和的模型    
with tf.Session() as sess:
    saver.restore(sess, "Saved_model/model.ckpt")
    print sess.run(result)

output：
[3.]

以上的保存和加载过程不同的就是，在加载模型的代码中没有运行变量的初始化过程，而是将变量的值通过已经爆粗你的模型加载进来。

2. 直接加载持久化的图

如果不希望对图上的运算进行重复定义，则可以直接加载已经持久化的图。

import tensorflow as tf

saver = tf.train.import_meta_graph("Saved_model/model.ckpt.meta")
with tf.Session() as sess:
    saver.restore(sess, "Saved_model/model.ckpt")
    print sess.run(tf.get_default_graph().get_tensor_by_name("add:0")) 

output:
[3.]

3. 变量重新命名

如果在保存或者在加载的时候给变量进行了重命名，则需要通过字典将模型保存时的变量名和需要加载的变量联系起来。

v1 = tf.Variable(tf.constant(1.0, shape=[1]), name = "other-v1")
v2 = tf.Variable(tf.constant(2.0, shape=[1]), name = "other-v2")
saver = tf.train.Saver({"v1": v1, "v2": v2})

4. 利用Saver类进行滑动平均的加载

在TensorFlow中，每一个变量的滑动平均值是通过影子变量维护的，所以要获取变量的滑动平均值实际上就是要获取这个影子变量的取值。如果在加载模型时，直接将影子变量映射到变量自身，那么在使用训练好的模型时，就不需要再调用函数来获取变量的滑动平均值了。

保存和加载滑动平均值

import tensorflow as tf

# 使用滑动平均
v = tf.Variable(0, dtype=tf.float32, name="v")
for variables in tf.global_variables(): print variables.name

ema = tf.train.ExponentialMovingAverage(0.99)
maintain_averages_op = ema.apply(tf.global_variables())
for variables in tf.global_variables(): print variables.name

output:
v:0
v:0
v/ExponentialMovingAverage:0

# 保存滑动平均模型
saver = tf.train.Saver()
with tf.Session() as sess:
    init_op = tf.global_variables_initializer()
    sess.run(init_op)

    sess.run(tf.assign(v, 10))
    sess.run(maintain_averages_op)
    # 保存的时候会将v:0  v/ExponentialMovingAverage:0这两个变量都存下来。
    saver.save(sess, "Saved_model/model2.ckpt")
    print sess.run([v, ema.average(v)])

output:
[10.0, 0.099999905]

# 加载滑动平均模型
v = tf.Variable(0, dtype=tf.float32, name="v")
# 通过变量重命名将原来变量v的滑动平均值直接赋值给v。
saver = tf.train.Saver({"v/ExponentialMovingAverage": v})
with tf.Session() as sess:
    saver.restore(sess, "Saved_model/model2.ckpt")
    print sess.run(v)

output:
0.0999999

为了使得更加方便地在加载时重命名滑动平均变量，tf.train.ExponentialMovingAverage类提供了variables_to_restore函数来生成tf.train.Saver类所需要的变量重命名字典。

import tensorflow as tf
v = tf.Variable(0, dtype=tf.float32, name="v")
ema = tf.train.ExponentialMovingAverage(0.99)
print ema.variables_to_restore()

saver = tf.train.Saver({"v/ExponentialMovingAverage": v})
with tf.Session() as sess:
    saver.restore(sess, "Saved_model/model2.ckpt")
    print sess.run(v)

output:
{u'v/ExponentialMovingAverage': 0x10f42a390>}
0.0999999

pb文件的保存和加载

使用convert_variables_to_constants函数，通过这个函数可以将计算图中的变量及其取值通过常量的方式来进行保存，这样整个TensorFlow中的计算图可以统一存放在一个文件中。

# pb文件的保存
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)
    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("Saved_model/combined_model.pb", "wb") as f:
           f.write(output_graph_def.SerializeToString())

output：
INFO:tensorflow:Froze 2 variables.
Converted 2 variables to const ops.

# pb文件的加载
from tensorflow.python.platform import gfile
with tf.Session() as sess:
    model_filename = "Saved_model/combined_model.pb"

    with gfile.FastGFile(model_filename, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())

    result = tf.import_graph_def(graph_def, return_elements=["add:0"])
    print sess.run(result)

output:
[array([ 3.], dtype=float32)]