Tensorflow custom saver that keeps VGG variables out of checkpoint files

custom_saver.py

class MySaver(tf.train.Saver):
    def __init__(self, var_list, extra_vars=None, extra_chkpt_file=None, **kwargs):
        super().__init__(var_list=var_list, **kwargs)
        self.extra_chkpt_file = extra_chkpt_file
        self.extra_saver = tf.train.Saver(var_list=extra_vars)
    
    def restore(self, sess, save_path):
        super().restore(sess, save_path)
        self.extra_saver.restore(sess, self.extra_chkpt_file)
                
# For training
variables_to_restore = slim.get_variables_to_restore(exclude=['vgg_19'])
vgg_vars = slim.get_model_variables('vgg_19')
saver = MySaver(var_list=variables_to_restore, extra_vars=vgg_vars, extra_chkpt_file=VGG_CHECKPOINT_FILE)
slim.learning.train(self.train_op, self.log_dir, saver=saver)