nolan-dev / GANInterface

Tool to interface with a StyleGAN model
Apache License 2.0
217 stars 25 forks source link

Conversion of .pkl to .pb file? #10

Open saadz-khan opened 4 years ago

saadz-khan commented 4 years ago

Actually I want to convert pkl to pb file with all the weights, variables and the tensorflow inference frozen graph stored into it? Can you even write the script for the conversion of such file? Will be of great help. The shell script docker run will convert all the parameters mentioned above? Some problems with permission denied with the shell script run.

nolan-dev commented 4 years ago

The script should include all the weights in the conversion. In my tests, inference using the frozen pb created by the script produced the same results as inference using the pkl. What was the error you encountered? It should hopefully still work in the gcp deep learning VM environment (https://cloud.google.com/deep-learning-vm) if you're having trouble in your own environment.

saadz-khan commented 4 years ago

I am getting the error command not found and sometimes cannot open the file here is my code. I dont know what is wrong with code followed all the steps listed in the read.txt of pkl_conversion. I am using google colab vm. `%cd /content/GANInterface/pkl_conversion/

!chmod +x '/content/GANInterface/pkl_conversion/download_and_patch.sh'

!chmod +x '/content/GANInterface/pkl_conversion/run_with_docker.sh'

!./run_with_docker.sh network-snapshot-010207.pkl 1024,1024`

nolan-dev commented 4 years ago

Is 'docker' the command that isn't found? Unless the image got changed all the other commands should be available in the docker container: my hope was that docker would be the only real requirement for running the conversion. I haven't used google colab so I'm not sure if it's available or easy to install there. It may be worth trying on a different machine that has docker installed.

saadz-khan commented 4 years ago

Hey nolan Can you help me with getting the code running. basically I want to get the stylegan model just like the tensorflow (.pb) or keras(h5) models which are present in the model libraries of different frameworks. A frozen graph with variables and weights. Did some research but can you help me getting this to load and is there any graph frozen into it? Will be a life saver if u could help me out.

import pickle
import numpy as np
import dnnlib.tflib as tflib
import dnnlib
import tensorflow as tf
import config

tflib.init_tf()
url = 'https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ'  # my trained data
with dnnlib.util.open_url(url, cache_dir=config.cache_dir) as f:
    _G, _D, Gs = pickle.load(f)
Gs.print_layers()
layers = Gs.list_layers()
latents_in = layers[0][1]
images_out = layers[-1][1]
sess = tf.get_default_session()
builder = tf.saved_model.builder.SavedModelBuilder('output_model')
builder.add_meta_graph_and_variables(sess, \
    [tf.saved_model.tag_constants.SERVING], \
        signature_def_map = {
            'serving_default': tf.saved_model.signature_def_utils.build_signature_def(
                inputs={'latents_in': tf.saved_model.utils.build_tensor_info(latents_in)},
                outputs={'images_out':  tf.saved_model.utils.build_tensor_info(images_out)},
                method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME,
            ),
        },
    strip_default_attrs=True)
builder.save()
nolan-dev commented 4 years ago

Hey sorry for the super-late reply. I tried my script on your pkl, and it seems like the issue is due to an incompatibility between the pkl and how I patch the network to save. I assumed completely default stylegan2 parameters, did you modify any hyperparameters before training? Unfortunately, the conversion script is super brittle