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, 7, 8, 30), 'tensorflow.InteractiveSession', 'tf.InteractiveSession', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((9, 11, 9, 28), 'pickle.load', 'pickle.load', ({(9, 23, 9, 27): 'file'}, {}), '(file)', False, 'import pickle\n'), ((10, 8, 10, 24), 'tensorflow.train.Saver', 'tf.train.Saver', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((15, 17, 15, 52), 'os.path.join', 'os.path.join', ({(15, 30, 15, 39): 'save_path', (15, 41, 15, 51): 'model_name'}, {}), '(save_path, model_name)', False, 'import os\n'), ((18, 7, 18, 47), 'tensorflow.train.get_checkpoint_state', 'tf.train.get_checkpoint_state', ({(18, 37, 18, 46): 'save_path'}, {}), '(save_path)', True, 'import tensorflow as tf\n'), ((19, 9, 19, 65), 'tensorflow.train.NewCheckpointReader', 'tf.train.NewCheckpointReader', ({(19, 38, 19, 64): 'ckpt.model_checkpoint_path'}, {}), '(ckpt.model_checkpoint_path)', True, 'import tensorflow as tf\n'), ((22, 0, 22, 51), 'wolframclient.serializers.export', 'wxf.export', (), '', True, 'import wolframclient.serializers as wxf\n'), ((13, 7, 13, 32), 'os.path.exists', 'os.path.exists', ({(13, 22, 13, 31): 'save_path'}, {}), '(save_path)', False, 'import os\n'), ((14, 4, 14, 26), 'os.makedirs', 'os.makedirs', ({(14, 16, 14, 25): 'save_path'}, {}), '(save_path)', False, 'import os\n'), ((25, 5, 25, 17), 'tensorflow.Session', 'tf.Session', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((27, 23, 32, 5), 'tensorflow.graph_util.convert_variables_to_constants', 'tf.graph_util.convert_variables_to_constants', (), '', True, 'import tensorflow as tf\n'), ((34, 9, 34, 57), 'tensorflow.gfile.GFile', 'tf.gfile.GFile', ({(34, 24, 34, 50): "('./target/' + name + '.pb')", (34, 52, 34, 56): '"""wb"""'}, {}), "('./target/' + name + '.pb', 'wb')", True, 'import tensorflow as tf\n'), ((26, 4, 26, 33), 'tensorflow.initialize_all_variables', 'tf.initialize_all_variables', ({}, {}), '()', True, 'import tensorflow as tf\n')]