Is it possible to make a trainable variable not trainable?
In order to remove a variable from the list of trainable variables, you can first access the collection through:
trainable_collection = tf.get_collection_ref(tf.GraphKeys.TRAINABLE_VARIABLES)
There, trainable_collection
contains a reference to the collection of trainable variables. If you pop elements from this list, doing for example trainable_collection.pop(0)
, you are going to remove the corresponding variable from the trainable variables, and thus this variable will not be trained.
Although this works with pop
, I am still struggling to find a way to correctly use remove
with the correct argument, so we don't depend on the index of the variables.
EDIT: Given that you have the name of the variables in the graph (you can obtain that by inspecting the graph protobuf or, what is easier, using Tensorboard), you can use it to loop through the list of trainable variables and then remove the variables from the trainable collection.
Example: say that I want the variables with names "batch_normalization/gamma:0"
and "batch_normalization/beta:0"
NOT to be trained, but they are already added to the TRAINABLE_VARIABLES
collection. What I can do is:
`
#gets a reference to the list containing the trainable variables
trainable_collection = tf.get_collection_ref(tf.GraphKeys.TRAINABLE_VARIABLES)
variables_to_remove = list()
for vari in trainable_collection:
#uses the attribute 'name' of the variable
if vari.name=="batch_normalization/gamma:0" or vari.name=="batch_normalization/beta:0":
variables_to_remove.append(vari)
for rem in variables_to_remove:
trainable_collection.remove(rem)
` This will successfully remove the two variables from the collection, and they will not be trained anymore.
When you want to train or optimize only certain layers of a pre-trained network, this is what you need to know.
TensorFlow's minimize
method takes an optional argument var_list
, a list of variables to be adjusted through back-propagation.
If you don't specify var_list
, any TF variable in the graph could be adjusted by the optimizer. When you specify some variables in var_list
, TF holds all other variables constant.
Here's an example of a script which jonbruner and his collaborator have used.
tvars = tf.trainable_variables()
g_vars = [var for var in tvars if 'g_' in var.name]
g_trainer = tf.train.AdamOptimizer(0.0001).minimize(g_loss, var_list=g_vars)
This finds all the variables they defined earlier that have "g_" in the variable name, puts them into a list, and runs the ADAM optimizer on them.
You can find the related answers here on Quora
After looking at the documentation and the code, I was not able to find a way to remove a Variable from the TRAINABLE_VARIABLES
.
Here is what happens:
- The first time
tf.get_variable('weights', trainable=True)
is called, the variable is added to the list ofTRAINABLE_VARIABLES
. - The second time you call
tf.get_variable('weights', trainable=False)
, you get the same variable but the argumenttrainable=False
has no effect as the variable is already present in the list ofTRAINABLE_VARIABLES
(and there is no way to remove it from there)
First solution
When calling the minimize
method of the optimizer (see doc.), you can pass a var_list=[...]
as argument with the variables you want to optimizer.
For instance, if you want to freeze all the layers of VGG except the last two, you can pass the weights of the last two layers in var_list
.
Second solution
You can use a tf.train.Saver()
to save variables and restore them later (see this tutorial).
- First you train your entire VGG model with all trainable variables. You save them in a checkpoint file by calling
saver.save(sess, "/path/to/dir/model.ckpt")
. - Then (in another file) you train the second version with non trainable variables. You load the variables previously stored with
saver.restore(sess, "/path/to/dir/model.ckpt")
.
Optionally, you can decide to save only some of the variables in your checkpoint file. See the doc for more info.