cvpr/how_to_extract_official_weights.md

Steps to extract official weights

Prerequisite

• install docker >= 19.03
• install nvidia-docker native

Step 1. Working on official repository

• Clone official repo

~$ git clone https://github.com/NVlabs/stylegan2.git
~$ cd stylegan2

• Download official weights (*.pkl)
• Create directory

~$ mkdir official_pretrained

• Place official weight (*.pkl) to <PATH-TO-OFFICIAL-REPO>/official_pretrained
• Create python script called export_weights.py and add following

import tensorflow as tf
import pretrained_networks
import pprint


network_pkl = './official_pretrained/stylegan2-ffhq-config-f.pkl'
print('Loading networks from "%s"...' % network_pkl)
_G, _D, Gs = pretrained_networks.load_networks(network_pkl)
print()

# Print network details.
print('===_G===')
_G.print_layers()
print('===Gs===')
Gs.print_layers()

t_var = tf.trainable_variables()
pprint.pprint(t_var)

saver = tf.train.Saver()

save_path = saver.save(tf.get_default_session(), "./model.ckpt")

• Create docker environment for stylegan2

~$ docker build -f Dockerfile -t stylegan2:official .

• Get inside docker container and run script

~$ docker run --gpus all --rm -u $(id -u):$(id -g) -it -v <PATH-TO-OFFICIAL-REPO>:/work-dir -w /work-dir stylegan2:official /bin/bash

# now you are inside docker container
# run following python script to extract model weights
~$ python export_weights.py

• Check that 4 files are created
  • checkpoint
  • model.ckpt.data-00000-of-00001
  • model.ckpt.index
  • model.ckpt.meta
• Leave docker container

~$ exit

Step 2. Working on this repository(stylegan2-tf-2.x)

• Move to stylegan2-tf-2.x directory

~$ cd <PATH-TO-stylegan2-tf-2.x>

• Copy 4 files that are created from step 1 to <PATH-TO-stylegan2-tf-2.x>/official-pretrained
• Create docker environment for stylegan2-tf-2.x

~$ docker build -f Dockerfile -t stylegan2:tf.2.x .

• Get inside docker container and run script

~$ docker run --gpus all --rm -u $(id -u):$(id -g) -it -v <PATH-TO-stylegan2-tf-2.x>:/work-dir -w /work-dir stylegan2:tf.2.x /bin/bash

# now you are inside docker container
# run following python script to convert model weights
~$ python inference_from_official_weights.py

• Check inference result images *.png
• Check converted model weights ./official-converted/cuda && ./official-converted/ref