详解TensorFlow的 tf.add_to_collection 函数:将变量添加到指定名称的集合中

  • Post category:Python

TensorFlow 中的 tf.add_to_collection 可以将某个张量添加到一个集合(collection)中。这个集合可以看成一个命名空间,用于保存张量或其他对象,方便以后使用或导出。tf.add_to_collection 函数的语法格式如下:

tf.add_to_collection(name, value)

其中,name 是集合的名称,value 是要添加到集合中的张量或其他对象,可以是一个 TensorFlow 的操作、一个张量、或一个 Python 对象。

tf.add_to_collection 函数可用于保存模型中的某些特殊操作,比如损失函数、权重等等。这些特殊操作可以通过 tf.get_collection 函数获取。tf.get_collection 函数的语法格式如下:

tf.get_collection(name, scope=None)

其中,name 是集合的名称,scope 是要获取的对象所在的命名空间,可以是字符串类型或 Scope 对象。若 scope 为 None,则获取指定名称的所有对象;若 scope 不为空,则获取指定名称下的所有对象。

下面给出两个例子。

示例一

在训练模型时,需要定义损失函数。为了方便管理,我们可以将不同的损失函数分别添加到名为 “losses” 的集合中,然后求和计算总损失。实现代码如下:

import tensorflow as tf

# 定义两个张量
w = tf.Variable(tf.random_normal([2, 1], stddev=1), name='weight')
b = tf.Variable(tf.zeros([1]), name='bias')

# 定义输入和输出节点
x = tf.placeholder(tf.float32, shape=[None, 2], name='input')
y = tf.placeholder(tf.float32, shape=[None, 1], name='output')

# 定义模型
z = tf.matmul(x, w) + b
output = tf.sigmoid(z)

# 定义损失函数
cross_entropy = -tf.reduce_mean(y * tf.log(output) + (1 - y) * tf.log(1 - output))

# 将损失函数添加到集合中
tf.add_to_collection('losses', cross_entropy)

# 计算总损失
loss = tf.add_n(tf.get_collection('losses'))

# 定义优化器
optimizer = tf.train.GradientDescentOptimizer(0.1)

# 定义训练过程
train_step = optimizer.minimize(loss)

# 运行训练过程
with tf.Session() as sess:
    tf.global_variables_initializer().run()
    # 执行多次训练过程
    for i in range(1000):
        sess.run(train_step, feed_dict={x: input_data, y: output_data})

示例二

在训练图像分类模型时,常常需要使用 L2 正则化来限制权重的大小,避免过拟合。为了方便管理,我们可以将所有需要 L2 正则化的权重分别添加到名为 “regularizer” 的集合中,然后统一计算正则化损失。实现代码如下:

import tensorflow as tf

# 定义两个张量
w1 = tf.Variable(tf.random_normal([2, 3], stddev=1), name='weight1')
w2 = tf.Variable(tf.random_normal([3, 1], stddev=1), name='weight2')
b = tf.Variable(tf.zeros([1]), name='bias')

# 定义输入和输出节点
x = tf.placeholder(tf.float32, shape=[None, 2], name='input')
y = tf.placeholder(tf.float32, shape=[None, 1], name='output')

# 定义模型
hidden = tf.nn.relu(tf.matmul(x, w1) + b)
output = tf.sigmoid(tf.matmul(hidden, w2) + b)

# 定义 L2 正则化损失
regularizer = tf.contrib.layers.l2_regularizer(scale=0.01)
tf.add_to_collection('regularizer', regularizer(w1))
tf.add_to_collection('regularizer', regularizer(w2))
reg_loss = tf.add_n(tf.get_collection('regularizer'))

# 定义总损失
cross_entropy = -tf.reduce_mean(y * tf.log(output) + (1 - y) * tf.log(1 - output))
loss = cross_entropy + reg_loss

# 定义优化器
optimizer = tf.train.GradientDescentOptimizer(0.1)

# 定义训练过程
train_step = optimizer.minimize(loss)

# 运行训练过程
with tf.Session() as sess:
    tf.global_variables_initializer().run()
    # 执行多次训练过程
    for i in range(1000):
        sess.run(train_step, feed_dict={x: input_data, y: output_data})