Tensorflow Graph and weights to json and back for training

tensorflow_json.py

import tensorflow as tf
import numpy as np
from google.protobuf import json_format
import json

np.random.seed(12345)


def tensorflow_get_weights():
    """
    @author https://github.com/maxim5 with code tweak for complete serialization
    """
    vs = tf.trainable_variables()
    values = tf.get_default_session().run(vs)
    return values


def tensorflow_set_weights(weights):
    """
    @author https://github.com/maxim5 with code tweak for complete serialization
    """
    assign_ops = []
    feed_dict = {}
    vs = tf.trainable_variables()
    zipped_values = zip(vs, weights)
    for var, value in zipped_values:
        value = np.asarray(value)
        assign_placeholder = tf.placeholder(var.dtype, shape=value.shape)
        assign_op = var.assign(assign_placeholder)
        assign_ops.append(assign_op)
        feed_dict[assign_placeholder] = value
    tf.get_default_session().run(assign_ops, feed_dict=feed_dict)


def convert_weights_to_json(weights):
    weights = [w.tolist() for w in weights]
    weights_list = json.dumps(weights)
    return weights_list


def convert_json_to_weights(json_weights):
    loaded_weights = json.loads(json_weights)
    loaded_weights = [np.asarray(x) for x in loaded_weights]
    return loaded_weights


def create_simple_graph():
    """
    Creates a very simple xor graph
    """
    x = tf.placeholder(tf.float32, shape=[None, 2], name='x')
    layer1 = tf.layers.dense(x, 12, activation=tf.nn.relu)
    layer2 = tf.layers.dense(layer1, 7, activation=tf.nn.relu)
    out = tf.layers.dense(layer2, 1, name='outer', activation=tf.nn.sigmoid)
    opt = tf.train.AdamOptimizer(learning_rate=.01)
    y = tf.placeholder(tf.float32, shape=[None, 1], name='y')

    loss = tf.reduce_mean(tf.square(y - out))
    mini = opt.minimize(loss, global_step=tf.train.get_or_create_global_step(), name='mini')
    return mini


def retrieve_xor():
    """
    Grabs xor data
    """
    xor = [(0.0, np.array([0.0, 0.0])),
           (0.0, np.array([1.0, 1.0])),
           (1.0, np.array([1.0, 0.0])),
           (1.0, np.array([0.0, 1.0]))]
    a = np.asarray([x for y, x in xor])
    b = np.asarray([y for y, _ in xor]).reshape((4, 1))
    return a, b


def run_initial_with_json_weights(opti, feed_dict):
    """
    returns both serialized json graph and weights
    """
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        for i in range(0, 100):
            sess.run(opti, feed_dict=feed_dict)
        first_weights = tensorflow_get_weights()
        g = tf.get_default_graph().as_graph_def()
    json_string = json_format.MessageToJson(g)
    return json_string, convert_weights_to_json(first_weights)


def run_serialized(json_graph, json_weights, feed_dict):
    """
    deserialize graph and run it again
    """
    gd = tf.GraphDef()
    gd = json_format.Parse(json_graph, gd)
    weights = convert_json_to_weights(json_weights)

    with tf.Session() as sess:
        tf.import_graph_def(gd)
        sess.run(tf.global_variables_initializer())
        nu_out = tf.get_default_graph().get_tensor_by_name('outer/Sigmoid:0')
        mini = tf.get_default_graph().get_tensor_by_name('mini:0')
        tensorflow_set_weights(weights)

        for i in range(0, 200):
            sess.run(mini, feed_dict=feed_dict)

        predicted = sess.run(nu_out, feed_dict=feed_dict)

    return predicted


def run_with_serialized_weights():
    """
    weights ARE turned into json
    """
    initial_graph = create_simple_graph()
    a,b = retrieve_xor()
    feed_dict = {'x:0': a, 'y:0': b}
    json_graph, json_weights = run_initial_with_json_weights(initial_graph, feed_dict)
    predictions = run_serialized(json_graph, json_weights, feed_dict)
    return predictions


if __name__ == "__main__":
    print(run_with_serialized_weights())