Running a pre-trained tensorflow StyleGan model in PyTorch


(Sam) #1

Is it possible to run a pre-trained Style model that was made in Tensorflow and have it be ported over to PyTorch to generate images?


#2

Yes! @tom and I transferred the pre-trained weights to PyTorch and created a notebook here. :slight_smile:


(Sam) #3

Thank you so much! So there should be no problem in using a pre-trained model that was made in TF (.pkl file) for this?


(Thomas V) #4

In the notebook, it says

For completeness, our conversion is below, but you can skip it if you download the PyTorch-ready weights.

so you can use weights from TF, but there are some conversion steps. @ptrblck and I put the steps we took for the original authors’ weights in the notebook to help you.

Best regards

Thomas


(Sam) #5

Gotcha. Thank you! I got it to work; however, the results are pretty lackluster. Should I play around with some of the parameters? (Note: I’m running on CPU mode, should I switch to a GPU ?)

Output below (I tried running it a few times and got similar blurry results) :


#6

Which weights are you using?


(Sam) #7

The ones in the reference code for now (same as the link above, Nvidia implementation):

karras2019stylegan-ffhq-1024x1024.pkl


#8

Are you running the notebook with these weights and get this output?
If so, did you change any hyperparameters or anything else in the models?
The outputs should look like the last two figures in the notebook.
Could you tell me which PyTorch version you are using so that we could debug it?


(Sam) #9

Yeah the exact same code as above minus the interpolation part at the very end.
My PyTorch version is on 1.1 (more exactly ‘1.1.0a0+fdaa77a’) and torchvision is on version 0.2.2


#10

Thanks, I’ll try to reproduce this issue.


(Sam) #11

Hmmm, I tried running it on Google Colab (PyTorch Version 1.0.1) with its 1 GPU and it seems to work there.
Not sure why it’s producing those blobs on my local (CPU only)