I have a pretrained model (checkpoint, tensorflow v1) with different variables and weights. I don't know all the variables, but I know two that I want to change their shape: v1 is in the shape of [4,768] and v2 is in the shape of [4]. I want to increase both to be [5,768] and [5] respectively and save the checkpoint again for fine-tuning purposes. To fill the missing data I want to take the average values of the variables.
Here is my code:
# The vars I want to change
v1 = tf.get_variable("v1", shape=[4, 768], initializer=utils.classification_initializer())
v2 = tf.get_variable("v2", shape=[4], initializer=tf.zeros_initializer())
checkpoint = {}
saver = tf.train.Saver()
with tf.Session() as sess:
# Restore checkpoint from source location (path).
saver.restore(sess, source)
# Get the vars values
checkpoint[v1.name] = v1.eval()
checkpoint[v2.name] = v2.eval()
new_data = {}
# Calc v1 average and reshape
avg = numpy.average(checkpoint[v1.name], axis=0)
new_data[v1.name] = numpy.vstack((checkpoint[v1.name], avg))
# Calc v2 average and reshape
avg = numpy.average(checkpoint[v2.name], axis=0)
new_data[v2.name] = numpy.append(checkpoint[v2.name], avg)
# Assign the new data and shape
sess.run(tf.assign(v1, new_data[v1.name], validate_shape=False))
sess.run(tf.assign(v2, new_data[v2.name], validate_shape=False))
# Save the checkpoint to target location (path).
saver.save(sess, target)
I was expecting to see a similar size model (the source checkpoint is about 1GB), but i get a much smaller file (target checkpoint is about 15KB). It seems that is saves only the variables that I've changed and not the entire checkpoint (other vars, weights, etc).
1 - this is the way to achieve my goal (reshaping and filling 2 vars in a checkpoint)?
2 - if so, how can i save the entire model (other vars, weights, etc) and not only the loaded vars?
from Changing the shape of a tensorflow variable in a checkpoint
No comments:
Post a Comment