tensorflow How to use TensorFlow Graph Collections? Create your own collection and use it to collect all your losses.


Example

Here we will create collection for losses of Neural Network's computational graph.

First create a computational graph like so:

with tf.variable_scope("Layer"):
    W = tf.get_variable("weights", [m, k],
        initializer=tf.zeros_initializer([m, k], dtype=tf.float32))
    b1 = tf.get_variable("bias", [k],
        initializer = tf.zeros_initializer([k], dtype=tf.float32))
    z = tf.sigmoid((tf.matmul(input, W) + b1))
    
    with tf.variable_scope("Softmax"):
        U = tf.get_variable("weights", [k, r],
            initializer=tf.zeros_initializer([k,r], dtype=tf.float32))
        b2 = tf.get_variable("bias", [r],
            initializer=tf.zeros_initializer([r], dtype=tf.float32))
    out = tf.matmul(z, U) + b2
cross_entropy = tf.reduce_mean(
    tf.nn.sparse_softmax_cross_entropy_with_logits(out, labels))

To create a new collection, you can simply start calling tf.add_to_collection() - the first call will create the collection.

tf.add_to_collection("my_losses", 
    self.config.l2 * (tf.add_n([tf.reduce_sum(U ** 2), tf.reduce_sum(W ** 2)])))
tf.add_to_collection("my_losses", cross_entropy)

And finally you can get tensors from your collection:

loss = sum(tf.get_collection("my_losses"))

Note that tf.get_collection() returns a copy of the collection or an empty list if the collection does not exist. Also, it does NOT create the collection if it does not exist. To do so, you could use tf.get_collection_ref() which returns a reference to the collection and actually creates an empty one if it does not exist yet.