How to restore the function defined in the graph?

Question

I defined a funciton in tensorflow as follows:

def generator(keep_prob, z, out_channel_dim, alphag1, is_train=True):
    """
    Create the generator network
    :param z: Input z
    :param out_channel_dim: The number of channels in the output image
    :param is_train: Boolean if generator is being used for training
    :return: The tensor output of the generator
    """
    # TODO: Implement Function
    # when it is training reuse=False
    # when it is not training reuse=True
    alpha=alphag1
    with tf.variable_scope('generator',reuse=not is_train):
        layer = tf.layers.dense(z, 3*3*512,activation=None,\
                                 kernel_initializer=tf.contrib.layers.xavier_initializer(uniform=False))
        layer = tf.reshape(layer, [-1, 3,3,512])
        layer = tf.layers.batch_normalization(layer, training=is_train)
        layer = tf.maximum(layer*alpha, layer)
        #layer = layer+tf.random_normal(shape=tf.shape(layer), mean=0.0, stddev=0.0001, dtype=tf.float32)
        #layer = tf.nn.dropout(layer,keep_prob)

        layer = tf.layers.conv2d_transpose(layer, 256, 4, strides=2, padding='same',\
                                 kernel_initializer=tf.contrib.layers.xavier_initializer_conv2d(uniform=False))
        layer = tf.layers.batch_normalization(layer, training=is_train)
        layer = tf.maximum(layer*alpha, layer)
        #layer = layer+tf.random_normal(shape=tf.shape(layer), mean=0.0, stddev=0.00001, dtype=tf.float32)
        #layer = tf.nn.dropout(layer,keep_prob)

        layer = tf.layers.conv2d_transpose(layer, 128, 4, strides=2, padding='same',\
                                 kernel_initializer=tf.contrib.layers.xavier_initializer_conv2d(uniform=False))
        layer = tf.layers.batch_normalization(layer, training=is_train)
        layer = tf.maximum(layer*alpha, layer)
        #layer = layer+tf.random_normal(shape=tf.shape(layer), mean=0.0, stddev=0.000001, dtype=tf.float32)
        #layer = tf.nn.dropout(layer,keep_prob)

        layer = tf.layers.conv2d_transpose(layer, 64, 4, strides=2, padding='same',\
                                 kernel_initializer=tf.contrib.layers.xavier_initializer_conv2d(uniform=False))
        layer = tf.layers.batch_normalization(layer, training=is_train)
        layer = tf.maximum(layer*alpha, layer)
        #layer = layer+tf.random_normal(shape=tf.shape(layer), mean=0.0, stddev=0.0000001, dtype=tf.float32)
        #layer = tf.nn.dropout(layer,keep_prob)

        layer = tf.layers.conv2d_transpose(layer, out_channel_dim, 4, strides=2, padding='same',\
                                 kernel_initializer=tf.contrib.layers.xavier_initializer_conv2d(uniform=False))
        #layer = layer+tf.random_normal(shape=tf.shape(layer), mean=0.0, stddev=0.00000001, dtype=tf.float32)

        layer = tf.tanh(layer)
    return layer

This is complicated such that to track each variable in each layer is difficult. I later used tf.train.Saver() and saver.save to save everything after training.

Now I would like to restore this function so that I can use it to do further manipulations while keeping the trained weigts of each layer unchanged.

I found online that most function like tf.get_default_graph().get_tensor_by_name or some other functions were limited to restore only the values of the variables but not this function.

For example the input z of this function generator(keep_prob, z, out_channel_dim, alphag1, is_train=True) is a tensor from another function. I want to restore this function so that I can use two new tensors z1 and z2 with she same shape as z.

layer1 = generator(keep_prob, z1, out_channel_dim, alphag1, is_train=False) 
layer2 = generator(keep_prob, z2, out_channel_dim, alphag1, is_train=False) 
layer = layer1 - layer2

and I can put this new tensor layer into another function. Here layer1 and layer2 use the function with the saved weights.

The thing that is difficlut is that when I use the function generator I have to specifiy it with the trianed weights which was stored using Saver(). I find it difficult to specify this function with its weights. For, 1. too many layers to track off and 2. I don't know how to specify weights for tf.layers.conv2().

So are there anyone who know how to solve this issue?

Answer 1

Why do you need to restore the function and what does that even mean? If you need to use a model, you have to restore the corresponding graph. What your function does is defining nodes of the graph. You may use your function to build or rebuild that graph again and then load weights stored somewhere using Saver() or you may restore graph from the protobuf file.

In order to rebuild the graph, try to invoke invoke your function somewhere output_layer=generator(keep_prob, z, out_channel_dim, alphag1, is_train=True) and than use Saver class as usual to restore weights. Your function does not compute, it defines a part or whole of the graph. All computations are performed by the graph.

In the last case you will find useful the following thread. Usually, you will need to know names of the input and output layers. That can be obtained by the code:

[n.name for n in tf.get_default_graph().as_graph_def().node]

Answer 2

After a long time of searching, it seems that maybe the following is a solution.

Define all the variables in advance,i.e.layer1 = generator(keep_prob, z1,

out_channel_dim, alphag1, is_train=False) 
layer2 = generator(keep_prob, z2, out_channel_dim, alphag1, is_train=False) 
layer = layer1 - layer2.

Now you can use tf.get_collection to find the operators.

It seems that tensorflow will not give you the pre defined functions. It keeps the graph and values only but not in the form of function. One needs to set everything needed in the furture in the graph or one should keep track of every weights, even too many.

Answer 3

This is a general question: I save the whole model into file and need to restore part of the model into part of new model.

Here name_map is a dict:the key is new name in graph and value is name in ckpt file.

def get_restore_saver(self, name_map, restore_optimise_var=True):

    var_grp = {_.op.name:_ for _ in tf.global_variables()}
    varm = {}
    for main_name in var_grp:
        if main_name in name_map:
            varm[name_map[main_name]] = var_grp[main_name]
        elif restore_optimise_var: # I use adam to optimise
            var_arr = main_name.split('/')
            tail = var_arr[-1]
            _ = '/'.join(var_arr[: -1])
            if tail in ['Adam', 'Adam_1', 'ExponentialMovingAverage'] and _ in name_map:
                varm[name_map[_] + "/" + tail] = var_grp[main_name]

    return tf.train.Saver(varm)