Loading model weights into a new tensorflow graph

Viewed 601

Goal

Using tensorflow, I'm trying to train a LSTM model for a certain number of iterations on data that's N timesteps per sample, then slowly increase the number of timesteps per sample as the model trains.

So maybe the RNN model is looking at 4 timesteps per training sample at first. After training for a while, performance levels out. I'd like to now continue training the model with 8 timesteps. This is basically a form of finetuning for RNNs.

Progress

The seemingly most straightforward way to do this would be to save the model after training it for a while, then rebuild a new graph with a new Variable X with more timesteps defined.

Unfortunately, I can't find a way to not hardcode the number of timesteps into my model. But that's ok, because if I recreate the model and fill it with saved weights, the shapes of the model shapes should be the same so it should work.

So I'm running the model a first time to generate a save file. Then I'm loading that save file and trying to populate a new graph with the weights from the old (almost identical) tensorflow graph.

This has been driving me crazy, so any help is much appreciated.

Code

Here's my code so far:

if MODEL_FILE is not None:
    # load from saved model file
    new_saver = tf.train.import_meta_graph(MODEL_FILE + '.meta')

weights = {
        'out': tf.Variable(tf.random_uniform([LSTM_SIZE, n_outputs_sm]))
        }
biases = {
        'out': tf.Variable(tf.random_uniform([n_outputs_sm]))
        }

# setup input X and output Y graph variables
x = tf.placeholder('float', [None, NUM_TIMESTEPS, n_input], name='input_x')
y = tf.placeholder('float', [None, n_outputs_sm], name='output_y')

# Feed forward function to get the RNN output. We're using a fancy type of LSTM cell.
def TFEncoderRNN(inp, weights, biases):
    # current_input_shape: (batch_size, n_steps, n_input
    # required shape: 'n_steps' tensors list of shape (batch_size, n_input)
    inp = tf.unstack(inp, NUM_TIMESTEPS, 1)
    lstm_cell = tf.contrib.rnn.LayerNormBasicLSTMCell(LSTM_SIZE, dropout_keep_prob=DROPOUT)
    outputs, states = tf.contrib.rnn.static_rnn(lstm_cell, inp, dtype=tf.float32)
    return tf.matmul(outputs[-1], weights['out']) + biases['out']

# we'll be able to call this to get our model output
pred = TFEncoderRNN(x, weights, biases)
# define loss and optimizer
cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=pred, labels=y))

# I define some more stuff here I'll leave out for brevity

init = None
if new_saver:
    new_saver.restore(sess, './' + MODEL_FILE)
    init = tf.initialize_variables([global_step])
else:
    init = tf.global_variables_initializer()
sess.run(init)

######
### TRAIN AND STUFF
######

print "Optimization finished!"

# save the current graph, you can just run this script again to
# continue training
if SAVE_MODEL:
    print "Saving model"
    saver = tf.train.Saver()
    saver.save(sess, 'tf_model_001')

Any ideas on how to move my trained model weights into a newly created graph/model?

1 Answers
Related