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', False, 'import pickle\n'), (10, 'tensorflow.train.Saver', 'tf.train.Saver', True, 'import tensorflow as tf\n'), (15, 'os.path.join', 'os.path.join', False, 'import os\n'), (18, 'tensorflow.train.get_checkpoint_state', 'tf.train.get_checkpoint_state', True, 'import tensorflow as tf\n'), (19, 'tensorflow.train.NewCheckpointReader', 'tf.train.NewCheckpointReader', True, 'import tensorflow as tf\n'), (22, 'wolframclient.serializers.export', 'wxf.export', True, 'import wolframclient.serializers as wxf\n'), (13, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (14, 'os.makedirs', 'os.makedirs', 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', True, 'import tensorflow as tf\n'), (34, 'tensorflow.gfile.GFile', 'tf.gfile.GFile', True, 'import tensorflow as tf\n'), (26, 'tensorflow.initialize_all_variables', 'tf.initialize_all_variables', True, 'import tensorflow as tf\n')]