Running a pre-trained tensorflow StyleGan model in PyTorch

sonamghosh: <p>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?</p>

ptrblck: <p>Yes! <a class="mention" href="/u/tom">@tom</a> and I transferred the pre-trained weights to PyTorch and created a notebook <a href="https://github.com/lernapparat/lernapparat/blob/master/style_gan/pytorch_style_gan.ipynb" rel="nofollow noopener">here</a>. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

sonamghosh: <p>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?</p>

tom: <p>In the notebook, it says</p>
<blockquote>
<p>For completeness, our conversion is below, but you can skip it if you download the PyTorch-ready weights.</p>
</blockquote>
<p>so you can use weights from TF, but there are some conversion steps. <a class="mention" href="/u/ptrblck">@ptrblck</a> and I put the steps we took for the original authors’ weights in the notebook to help you.</p>
<p>Best regards</p>
<p>Thomas</p>

sonamghosh: <p>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 ?)</p>
<p>Output below (I tried running it a few times and got similar blurry results) :<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/b/b54617823789a80c9e1256abbc33996704cf5d8e.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/b54617823789a80c9e1256abbc33996704cf5d8e" title="50%20PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/b/b54617823789a80c9e1256abbc33996704cf5d8e_2_690x279.jpeg" alt="50%20PM" data-base62-sha1="pRCAMhAQBGr1QovOntn7Ut8CSTQ" width="690" height="279" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/b/b54617823789a80c9e1256abbc33996704cf5d8e_2_690x279.jpeg, https://discuss.pytorch.org/uploads/default/optimized/2X/b/b54617823789a80c9e1256abbc33996704cf5d8e_2_1035x418.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/2X/b/b54617823789a80c9e1256abbc33996704cf5d8e_2_1380x558.jpeg 2x" data-dominant-color="B4A79F"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">50%20PM</span><span class="informations">3002×1214 1.76 MB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

ptrblck: <p>Which weights are you using?</p>

sonamghosh: <p>The ones in the reference code for now (same as the link above, Nvidia implementation):</p>
<p>karras2019stylegan-ffhq-1024x1024.pkl</p>

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

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

ptrblck: <p>Thanks, I’ll try to reproduce this issue.</p>

sonamghosh: <p>Hmmm, I tried running it on Google Colab (PyTorch Version 1.0.1) with its 1 GPU and it seems to work there.<br>
Not sure why it’s producing those blobs on my local (CPU only)<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/a/aa02a530585b89d2b846c9b05db68119a695598f.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/aa02a530585b89d2b846c9b05db68119a695598f" title="13%20PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/a/aa02a530585b89d2b846c9b05db68119a695598f_2_690x434.jpeg" alt="13%20PM" data-base62-sha1="ofYPV78ApeU6alQxsQyJxzyhARN" width="690" height="434" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/a/aa02a530585b89d2b846c9b05db68119a695598f_2_690x434.jpeg, https://discuss.pytorch.org/uploads/default/optimized/2X/a/aa02a530585b89d2b846c9b05db68119a695598f_2_1035x651.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/2X/a/aa02a530585b89d2b846c9b05db68119a695598f_2_1380x868.jpeg 2x" data-dominant-color="6F635B"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">13%20PM</span><span class="informations">1558×982 372 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

bgavran: <p>Thank you so much - this is very well explained and easy to use. Correct me if I’m wrong - but the weights have been converted for the generator only and not for the discriminator?</p>

ptrblck: <p>Yes, the notebook doesn’t contain the discriminator implementation.</p>

tom: <p>Well, I did eventually convert the discriminator, too, but never got around to putting it up. If it’s of help, I can put up my (very raw) notebook for that, too.</p>

bgavran: <p>Yes, that’d be great!</p>

tom: <p><a href="https://github.com/lernapparat/lernapparat/blob/master/style_gan/pytorch_style_gan_test_discriminator.ipynb" rel="nofollow noopener">Here</a>, but it’s literally just the discriminator model.</p>

kylepob: <p>Tom,</p>
<p>First off, great work, and thank you for sharing. I’ve trained a transfer learned model from the NVIDIA FFHQ model in TensorFlow. Using your code to load this transfer learned model, it produces the appropriate images, but the images have a muted dynamic range/strange color space. Effectively, they look exactly like the images the NVIDIA model produces, except there is a pinkish/gray tinge to them. Therefore, the network output is not the same between the two models. Unfortunately, I can’t share examples of the phenomena.</p>

tom: <p>It sounds a bit like you have different parameters. We specialized to / hardcoded the parameters of the default model. You’d be able to modify the function to better reflect your output characteristics.</p>
<p>Best regards</p>
<p>Thomas</p>

kylepob: <p>Thanks for the reply Tom. Do you know specifically where I should start with these modifications? I believed I was using the default configuration of the StyleGAN repo, and assumed that only the parameters values would change, after transfer learning (rather than the the operations or the weight/bias dimensions). One other issue I was concerned about was this bug in the StyleGAN repo.</p>
<aside class="onebox githubpullrequest">
  <header class="source">
      <a href="https://github.com/NVlabs/stylegan/pull/3" target="_blank" rel="nofollow noopener">github.com/NVlabs/stylegan</a>
  </header>
  <article class="onebox-body">
      <a href="https://github.com/khalilsarwari" target="_blank" rel="nofollow noopener">
    <img alt="khalilsarwari" src="https://avatars0.githubusercontent.com/u/20780564?v=4" class="thumbnail onebox-avatar" width="90" height="90">
  </a>

<h4>
  <a href="https://github.com/NVlabs/stylegan/pull/3" target="_blank" rel="nofollow noopener">Dynamic Range Mapping Bug</a>
</h4>

<div class="date">
  by <a href="https://github.com/khalilsarwari" target="_blank" rel="nofollow noopener">khalilsarwari</a>
  on <a href="https://github.com/NVlabs/stylegan/pull/3" target="_blank" rel="nofollow noopener">09:30AM - 25 Mar 19 UTC</a>
</div>

<div class="github-commit-stats">
  <strong>1 commits</strong>
  changed <strong>1 files</strong>
  with <strong>1 additions</strong>
  and <strong>1 deletions</strong>.
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>Is it possible this would impact how the image is postprocessed?</p>

kylepob: <p>Just in case anyone runs into the issue I described. The error was related to the fmap_decay parameter. I was trying to convert a StyleGAN-Tensorflow trained model that had been checkpointed halfway through the training of the 1024x1024 LOD (level of detail). As a result the alpha/fmap_decay  parameter for the upscaling layer was only @~0.5. That led to me getting strange images, because the code posted here assumes an fmap_decay of 1.0 (which is appropriate for a fully trained LOD). If you plan on using a model that isn’t fully trained/checkpointed at the end of an LOD, try tuning this parameter so that you can get the same quality of image that you are seeing output by the tensorflow model.</p>