import os
import pickle
import tensorflow as tf
import wolframclient.serializers as wxf

name = 'karras2018iclr-celebahq-1024x1024'
file = open(name + '.pkl', 'rb')
sess = tf.InteractiveSession()
G, D, Gs = pickle.load(file)
saver = tf.train.Saver()
save_path = "./target/" + name + "/"
model_name = 'model'
if not os.path.exists(save_path):
    os.makedirs(save_path)
save_path_full = os.path.join(save_path, model_name)
saver.save(sess, save_path_full)

ckpt = tf.train.get_checkpoint_state(save_path)
reader = tf.train.NewCheckpointReader(ckpt.model_checkpoint_path)
all_variables = list(reader.get_variable_to_shape_map().keys())
npy = dict(zip(all_variables, map(reader.get_tensor, all_variables)))
wxf.export(npy, name + '.wxf', target_format='wxf')

# Save as protobuf
with tf.Session() as sess:
    tf.initialize_all_variables().run()
    output_graph_def = tf.graph_util.convert_variables_to_constants(
        sess=sess,
        input_graph_def=sess.graph_def,
        # output_node_names=['G_paper_1/images_out']
        output_node_names=['G_paper_1/ToRGB_lod0/add']
    )

    with tf.gfile.GFile("./target/" + name + ".pb", "wb") as file:  # 保存模型
        file.write(output_graph_def.SerializeToString())  # 序列化输出

[(8, 'tensorflow.InteractiveSession', 'tf.InteractiveSession', ([], {}), True, 'import tensorflow as tf\n'), (9, 'pickle.load', 'pickle.load', (['file'], {}), False, 'import pickle\n'), (10, 'tensorflow.train.Saver', 'tf.train.Saver', ([], {}), True, 'import tensorflow as tf\n'), (15, 'os.path.join', 'os.path.join', (['save_path', 'model_name'], {}), False, 'import os\n'), (18, 'tensorflow.train.get_checkpoint_state', 'tf.train.get_checkpoint_state', (['save_path'], {}), True, 'import tensorflow as tf\n'), (19, 'tensorflow.train.NewCheckpointReader', 'tf.train.NewCheckpointReader', (['ckpt.model_checkpoint_path'], {}), True, 'import tensorflow as tf\n'), (22, 'wolframclient.serializers.export', 'wxf.export', (['npy', "(name + '.wxf')"], {'target_format': '"""wxf"""'}), True, 'import wolframclient.serializers as wxf\n'), (13, 'os.path.exists', 'os.path.exists', (['save_path'], {}), False, 'import os\n'), (14, 'os.makedirs', 'os.makedirs', (['save_path'], {}), False, 'import os\n'), (25, 'tensorflow.Session', 'tf.Session', ([], {}), True, 'import tensorflow as tf\n'), (27, 'tensorflow.graph_util.convert_variables_to_constants', 'tf.graph_util.convert_variables_to_constants', ([], {'sess': 'sess', 'input_graph_def': 'sess.graph_def', 'output_node_names': "['G_paper_1/ToRGB_lod0/add']"}), True, 'import tensorflow as tf\n'), (34, 'tensorflow.gfile.GFile', 'tf.gfile.GFile', (["('./target/' + name + '.pb')", '"""wb"""'], {}), True, 'import tensorflow as tf\n'), (26, 'tensorflow.initialize_all_variables', 'tf.initialize_all_variables', ([], {}), True, 'import tensorflow as tf\n')]