Tuesday, 24 August 2021

Changing the shape of a tensorflow variable in a checkpoint

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