TensorFlow: resume training with different loss function

Question

This question is of similar spirit as this one, but I simply can't get it to work, so I'll need some help...

Aim: To pre-train a network using TensorFlow with a rather crude loss function, so that all weights/biases are in the right ball park. Then, continue training with a more sophisticated loss function to fine-tune the model.

Attempted approach: After the pre-training stage, save the model (using tf.train.Saver()), then immediately re-load the variables and change loss function. Below I've added a minimalistic non-working example:

# Crude loss function
loss_fn_1 = tf.reduce_mean(tf.losses.huber_loss(labels=box_in, predictions=box_out), name="loss")

# Less crude loss function
loss_fn_2 = tf.reduce_mean(model.IOU_loss(box_in, box_out), name="IOU_loss")

train_step = tf.train.AdamOptimizer(1e-4).minimize(loss_fn_1)
saver = tf.train.Saver()

# First round of training
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    # Train hard, fight easy
    saver.save(sess, savefile)

# Second round of training
with tf.Session() as sess:

    saver = tf.train.import_meta_graph(savefile_meta)
    saver.restore(sess, tf.train.latest_checkpoint(savedir))

    graph = tf.get_default_graph()
    IOU_loss = graph.get_tensor_by_name("IOU_loss:0")

    train_step = tf.train.AdamOptimizer(1e-4).minimize(IOU_loss)

    # 'Eye of the tiger'-type of training

I've tried various combinations of re-defining the train_step before or after re-initiating the tf.Session(), and loading/renaming various other variables. The main problem is that I'm pretty clueless of what I'm doing, or what the errors really mean, and my random walk hasn't gotten me anywhere.

MPA · Accepted Answer · 2018-01-03T15:55:26.100

After random walking for a bit longer, I found the following solution that worked for me: you can simply define a second optimisation function - see example below

# Crude loss function
loss_fn_1 = tf.reduce_mean(tf.losses.huber_loss(labels=box_in, predictions=box_out), name="loss")

# Less crude loss function
loss_fn_2 = tf.reduce_mean(model.IOU_loss(box_in, box_out), name="IOU_loss")

trainer1 = tf.train.AdamOptimizer(1e-4).minimize(loss_fn_1)
trainer2 = tf.train.AdamOptimizer(1e-4).minimize(loss_fn_2)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())

    # First round of training
    for i in range(100):
        trainer1.run()

    # Second round of training
    for i in range(1000):
        trainer2.run()

Please comment/answer if you feel other solutions are preferable.

TensorFlow: resume training with different loss function

1 Answers1