mangle_tf_checkpoint.py

import os
import sys

import tensorflow as tf

logdir = "/tmp/checkpoints"


def get_variables(logdir):
    variables = {}
    for name, shape in tf.contrib.framework.list_variables(logdir):
        variables[name] = tf.contrib.framework.load_variable(logdir, name)
    return variables


def main():
    variables = get_variables(logdir)

    mode = sys.argv[1]
    if mode == "print":
        for name, var in variables.items():
            print(name, var)

    if mode == "change":
        checkpoint = tf.train.get_checkpoint_state(logdir)
        with tf.Session() as sess:
            for name, var in variables.items():
                var.fill(42.0)
                tfvar = tf.Variable(var, name=name)
            saver = tf.train.Saver()
            sess.run(tf.global_variables_initializer())
            saver.save(sess, checkpoint.model_checkpoint_path)


if __name__ == "__main__":
    main()