init

.gitignore

@@ -0,0 +1,135 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pycharm
+.idea/
+
+# vscode
+.vscode/

README.md

@@ -1 +1,49 @@
-# StyleGAN2_tf2_conditioned
+# stylegan2-tf-2.x
+* This is tensorflow 2.0 based keras subclassing reimplementation of official [StyleGAN2 Repo](https://github.com/NVlabs/stylegan2)
+
+## Inference from official weights
+* Please check [how_to_extract_official_weights.md](./how_to_extract_official_weights.md)
+
+| Official repo result | Official weight copied result from this repo |
+| :---: | :---: |
+| ![official_result] | ![restored_result] |
+
+## Training on 256x256
+* batch size 32
+* around 300k steps
+
+| d_loss | g_loss |
+| :---: | :---: |
+| ![256x256_d_loss]| ![256x256_g_loss]|
+
+* result samples
+
+| | | |
+| :---: | :---: | :---: |
+| ![256x256_result_0]| ![256x256_result_1]| ![256x256_result_2]|
+
+## Using custom data
+* Create dataset(tf.data.Dataset) loader which returns 
+  * shape: `[None, 3, res, res]` where `res` is desired image resolution
+  * image value range: `[-1, +1]`
+ 
+## Etc
+### Using with pycharm
+* If .bashrc file's paths don't work in pycharm environment (e.g. `nvcc -h`), open pycharm with following
+```bash
+# move to pycharm installed location (location may vary)
+moono@moono-ubuntu:~$ cd .local/share/JetBrains/Toolbox/apps/PyCharm-P/ch-0/201.7846.77/bin/
+
+# launch pycharm with .bashrc 
+moono@moono-ubuntu:~/.local/share/JetBrains/Toolbox/apps/PyCharm-P/ch-0/201.7846.77/bin$ bash pycharm.sh
+```
+
+[loss_tensorboard]: assets/tf-keras-stylegan2-loss.PNG
+[generation_tensorboard]: assets/tf-keras-stylegan2-fake-images.PNG
+[official_result]: assets/seed6600-official.png
+[restored_result]: assets/seed6600-restored.png
+[256x256_d_loss]: assets/d_loss_256x256.png
+[256x256_g_loss]: assets/g_loss_256x256.png
+[256x256_result_0]: assets/out_256x256_0.png
+[256x256_result_1]: assets/out_256x256_1.png
+[256x256_result_2]: assets/out_256x256_2.png

copy_official_weights.py

@@ -0,0 +1,277 @@
+import os
+import tensorflow as tf
+
+from load_models import load_generator, load_discriminator
+
+
+def handle_mapping(w_name, is_g_clone):
+    def extract_info(name):
+        splitted = name.split('/')
+        index = splitted.index('g_mapping')
+        indicator = splitted[index + 1]
+        val = indicator.split('_')[-1]
+        return val
+
+    level = extract_info(w_name)
+    o_prefix = f'G_mapping_1' if is_g_clone else f'G_mapping'
+    if 'w' in w_name:
+        official_var_name = f'{o_prefix}/Dense{level}/weight'
+    else:
+        official_var_name = f'{o_prefix}/Dense{level}/bias'
+    return official_var_name
+
+
+def to_rgb_layer(name, r, o_prefix):
+    if 'conv/w:0' in name:
+        o_name = f'{o_prefix}/{r}x{r}/ToRGB/weight'
+    elif 'mod_dense/w:0' in name:
+        o_name = f'{o_prefix}/{r}x{r}/ToRGB/mod_weight'
+    elif 'mod_bias/b:0' in name:
+        o_name = f'{o_prefix}/{r}x{r}/ToRGB/mod_bias'
+    else:   # if 'bias/b:0' in name:
+        o_name = f'{o_prefix}/{r}x{r}/ToRGB/bias'
+    return o_name
+
+
+def handle_block_layer(name, r, o_prefix):
+    if 'conv_0/w:0' in name:
+        o_name = f'{o_prefix}/{r}x{r}/Conv0_up/weight'
+    elif 'conv_0/mod_dense/w:0' in name:
+        o_name = f'{o_prefix}/{r}x{r}/Conv0_up/mod_weight'
+    elif 'conv_0/mod_bias/b:0' in name:
+        o_name = f'{o_prefix}/{r}x{r}/Conv0_up/mod_bias'
+    elif 'noise_0/w:0' in name:
+        o_name = f'{o_prefix}/{r}x{r}/Conv0_up/noise_strength'
+    elif 'bias_0/b:0' in name:
+        o_name = f'{o_prefix}/{r}x{r}/Conv0_up/bias'
+    elif 'conv_1/w:0' in name:
+        o_name = f'{o_prefix}/{r}x{r}/Conv1/weight'
+    elif 'conv_1/mod_dense/w:0' in name:
+        o_name = f'{o_prefix}/{r}x{r}/Conv1/mod_weight'
+    elif 'conv_1/mod_bias/b:0' in name:
+        o_name = f'{o_prefix}/{r}x{r}/Conv1/mod_bias'
+    elif 'noise_1/w:0' in name:
+        o_name = f'{o_prefix}/{r}x{r}/Conv1/noise_strength'
+    else:   # if 'bias_1/b:0' in name:
+        o_name = f'{o_prefix}/{r}x{r}/Conv1/bias'
+    return o_name
+
+
+def handle_const_layer(name, o_prefix):
+    if 'const:0' in name:
+        o_name = f'{o_prefix}/4x4/Const/const'
+    elif 'conv/w:0' in name:
+        o_name = f'{o_prefix}/4x4/Conv/weight'
+    elif 'mod_dense/w:0' in name:
+        o_name = f'{o_prefix}/4x4/Conv/mod_weight'
+    elif 'mod_bias/b:0' in name:
+        o_name = f'{o_prefix}/4x4/Conv/mod_bias'
+    elif 'noise/w:0' in name:
+        o_name = f'{o_prefix}/4x4/Conv/noise_strength'
+    else:
+        o_name = f'{o_prefix}/4x4/Conv/bias'
+    return o_name
+    
+    
+def handle_synthesis(w_name, is_g_clone):
+    def extract_info(name):
+        splitted = name.split('/')
+        index = splitted.index('g_synthesis')
+        indicator1 = splitted[index + 1]
+        indicator2 = splitted[index + 2]
+        r = indicator1.split('x')[1]
+        d = indicator2
+        return r, d
+
+    res, divider = extract_info(w_name)
+    o_prefix = f'G_synthesis_1' if is_g_clone else f'G_synthesis'
+    if divider == 'ToRGB':
+        official_var_name = to_rgb_layer(w_name, res, o_prefix)
+    elif divider == 'block':
+        official_var_name = handle_block_layer(w_name, res, o_prefix)
+    else:   # const
+        official_var_name = handle_const_layer(w_name, o_prefix)
+    return official_var_name
+
+
+def handle_discriminator_layer(w_name):
+    def extract_info(name):
+        splitted = name.split('/')
+        resolution = splitted[1]
+        resolution = resolution.split('x')[0]
+        return resolution
+
+    res = extract_info(w_name)
+    if 'last_dense' in w_name:
+        o_name = 'D/Output/weight'
+    elif 'last_bias' in w_name:
+        o_name = 'D/Output/bias'
+    elif 'FromRGB' in w_name:
+        o_name = 'D/1024x1024/FromRGB/weight' if 'conv_' in w_name else 'D/1024x1024/FromRGB/bias'
+    elif 'skip' in w_name:
+        o_name = f'D/{res}x{res}/Skip/weight'
+    elif 'dense_' in w_name:
+        o_name = f'D/4x4/Dense0/weight'
+    elif 'conv_' in w_name:
+        if res == '4':
+            o_name = f'D/{res}x{res}/Conv/weight'
+        else:
+            o_name = f'D/{res}x{res}/Conv0/weight' if '_0' in w_name else f'D/{res}x{res}/Conv1_down/weight'
+    elif 'bias_' in w_name:
+        if res == '4':
+            o_name = f'D/{res}x{res}/Conv/bias' if '_0' in w_name else f'D/{res}x{res}/Dense0/bias'
+        else:
+            o_name = f'D/{res}x{res}/Conv0/bias' if '_0' in w_name else f'D/{res}x{res}/Conv1_down/bias'
+    else:
+        raise ValueError('Something went wrong!!')
+    return o_name
+
+
+def variable_name_mapper_g(g, is_g_clone):
+    name_mapper = dict()
+    for w in g.weights:
+        w_name, w_shape = w.name, w.shape
+
+        # mapping layer
+        if 'g_mapping' in w_name:
+            official_var_name = handle_mapping(w_name, is_g_clone)
+        elif 'g_synthesis' in w_name:
+            official_var_name = handle_synthesis(w_name, is_g_clone)
+        else:
+            # w_avg
+            official_var_name = 'Gs/dlatent_avg' if is_g_clone else 'G/dlatent_avg'
+
+        name_mapper[official_var_name] = w
+    return name_mapper
+
+
+def variable_name_mapper_d(d):
+    name_mapper = dict()
+    for w in d.weights:
+        w_name, w_shape = w.name, w.shape
+
+        official_var_name = handle_discriminator_layer(w_name)
+        name_mapper[official_var_name] = w
+    return name_mapper
+
+
+def check_shape(name_mapper, official_vars):
+    for official_name, v in name_mapper.items():
+        official_shape = [s for n, s in official_vars if n == official_name][0]
+
+        if official_shape == v.shape:
+            print('{}: shape matches'.format(official_name))
+        else:
+            # print(f'Official: {official_name} -> {official_shape}')
+            # print(f'Current: {v.name} -> {v.shape}')
+            raise ValueError('{}: wrong shape'.format(official_name))
+    return
+
+
+def convert_official_generator_weights(ckpt_dir, is_g_clone, use_custom_cuda):
+    generator = load_generator(g_params=None, is_g_clone=is_g_clone, ckpt_dir=None, custom_cuda=use_custom_cuda)
+
+    # restore official ones to current implementation
+    official_checkpoint = tf.train.latest_checkpoint('./official-pretrained')
+    official_vars = tf.train.list_variables(official_checkpoint)
+
+    # get name mapper
+    name_mapper = variable_name_mapper_g(generator, is_g_clone=is_g_clone)
+    for name_g, tvar in name_mapper.items():
+        print(f'{name_g}: {tvar.name}')
+
+    # check shape
+    check_shape(name_mapper, official_vars)
+
+    # restore
+    tf.compat.v1.train.init_from_checkpoint(official_checkpoint, assignment_map=name_mapper)
+
+    # save
+    if is_g_clone:
+        ckpt = tf.train.Checkpoint(g_clone=generator)
+    else:
+        ckpt = tf.train.Checkpoint(generator=generator)
+    out_dir = os.path.join(ckpt_dir, 'g_clone' if is_g_clone else 'generator')
+    manager = tf.train.CheckpointManager(ckpt, out_dir, max_to_keep=1)
+    manager.save(checkpoint_number=0)
+    return
+
+
+def convert_official_discriminator_weights(ckpt_dir, use_custom_cuda):
+    discriminator = load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=use_custom_cuda)
+
+    # restore official ones
+    official_checkpoint = tf.train.latest_checkpoint('./official-pretrained')
+    official_vars = tf.train.list_variables(official_checkpoint)
+
+    # get name mapper
+    name_mapper = variable_name_mapper_d(discriminator)
+    for name_d, tvar in name_mapper.items():
+        print(f'{name_d}: {tvar.name}')
+
+    # check shape
+    check_shape(name_mapper, official_vars)
+
+    # restore
+    tf.compat.v1.train.init_from_checkpoint(official_checkpoint, assignment_map=name_mapper)
+
+    # save
+    ckpt = tf.train.Checkpoint(discriminator=discriminator)
+    out_dir = os.path.join(ckpt_dir, 'discriminator')
+    manager = tf.train.CheckpointManager(ckpt, out_dir, max_to_keep=1)
+    manager.save(checkpoint_number=0)
+    return
+
+
+def convert_official_weights_together(ckpt_dir, use_custom_cuda):
+    # instantiate all models
+    discriminator = load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=use_custom_cuda)
+    generator = load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=use_custom_cuda)
+    g_clone = load_generator(g_params=None, is_g_clone=True, ckpt_dir=None, custom_cuda=use_custom_cuda)
+
+    # restore official ones
+    official_checkpoint = tf.train.latest_checkpoint('./official-pretrained')
+    official_vars = tf.train.list_variables(official_checkpoint)
+    for name, shape in official_vars:
+        print(f'{name}: {shape}')
+
+    # get name mapper
+    name_mapper_d = variable_name_mapper_d(discriminator)
+    name_mapper_g1 = variable_name_mapper_g(generator, is_g_clone=False)
+    name_mapper_g2 = variable_name_mapper_g(g_clone, is_g_clone=True)
+    name_mapper = {**name_mapper_d, **name_mapper_g1, **name_mapper_g2}
+
+    # check shape
+    check_shape(name_mapper, official_vars)
+
+    # restore
+    tf.compat.v1.train.init_from_checkpoint(official_checkpoint, assignment_map=name_mapper)
+
+    # save
+    ckpt = tf.train.Checkpoint(discriminator=discriminator,
+                               generator=generator,
+                               g_clone=g_clone)
+    manager = tf.train.CheckpointManager(ckpt, ckpt_dir, max_to_keep=1)
+    manager.save(checkpoint_number=0)
+    return
+
+
+def main():
+    from tf_utils import allow_memory_growth
+
+    allow_memory_growth()
+
+    ckpt_dir_base = './official-converted'
+    for use_custom_cuda in [True, False]:
+        ckpt_dir = os.path.join(ckpt_dir_base, 'cuda') if use_custom_cuda else os.path.join(ckpt_dir_base, 'ref')
+
+        convert_official_weights_together(ckpt_dir, use_custom_cuda)
+
+        # convert_official_discriminator_weights(ckpt_dir, use_custom_cuda)
+        # for is_g_clone in [True, False]:
+        #     convert_official_generator_weights(ckpt_dir, is_g_clone, use_custom_cuda)
+    return
+
+
+if __name__ == '__main__':
+    main()

dataset.py

@@ -0,0 +1,41 @@
+import os
+import sys
+import time
+import cv2
+import numpy as np
+import tensorflow as tf
+
+class DatasetSimple(tf.data.Dataset):
+
+    def _generator(folder):
+        # Opening the file
+        folder = folder.decode("utf-8")
+        uv_folder = os.path.join(folder, 'npz') 
+        pose_hape_folder = os.path.join(folder, 'poses_shapes') 
+        frames = os.listdir(uv_folder)
+
+        for frame in frames:
+            uv = np.load(os.path.join(uv_folder, frame))
+            uv = uv['uv_gt']
+            uv = cv2.resize(uv, [64, 128])
+            uv = np.transpose(uv, [2, 0, 1])
+
+            pose_shape = np.load(os.path.join(pose_hape_folder, frame.split('.')[0] + '.npy'))
+
+            yield uv, pose_shape
+
+    def __new__(cls, folder, res=64):
+        return tf.data.Dataset.from_generator(
+            cls._generator,
+            output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (83), dtype = tf.float32)),
+            args=(folder, )
+        )
+
+if __name__ == '__main__':
+    
+    dataset = DatasetSimple('/workspace/Dataset/stylegan/npz_1')
+
+    for d in dataset:
+        d
+        pass
+    

dataset_ffhq.py

@@ -0,0 +1,62 @@
+import os
+import numpy as np
+import tensorflow as tf
+
+
+# n_samples = 70000
+def parse_tfrecord_tf(record):
+    features = tf.io.parse_single_example(record, features={
+        'shape': tf.io.FixedLenFeature([3], tf.dtypes.int64),
+        'data': tf.io.FixedLenFeature([], tf.dtypes.string)
+    })
+
+    # [0 ~ 255] uint8 -> [-1.0 ~ 1.0] float32
+    image = tf.io.decode_raw(features['data'], tf.dtypes.uint8)
+    image = tf.reshape(image, features['shape'])
+    image = tf.transpose(image, perm=[1, 2, 0])
+    image = tf.image.random_flip_left_right(image)
+    image = tf.cast(image, tf.dtypes.float32)
+    image = image / 127.5 - 1.0
+    image = tf.transpose(image, perm=[2, 0, 1])
+    return image
+
+
+def get_ffhq_dataset(tfrecord_base_dir, res, batch_size, epochs=None, buffer_size=1000):
+    fn_index = int(np.log2(res))
+    tfrecord_fn = os.path.join(tfrecord_base_dir, 'ffhq-r{:02d}.tfrecords'.format(fn_index))
+
+    # with tf.device('/cpu:0'):
+    dataset = tf.data.TFRecordDataset(tfrecord_fn)
+    dataset = dataset.map(map_func=parse_tfrecord_tf, num_parallel_calls=8)
+    dataset = dataset.shuffle(buffer_size=buffer_size)
+    dataset = dataset.repeat(epochs)
+    dataset = dataset.batch(batch_size)
+    dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
+    return dataset
+
+
+def main():
+    from PIL import Image
+
+    res = 64
+    batch_size = 4
+    epochs = 1
+    tfrecord_dir = './tfrecords'
+
+    dataset = get_ffhq_dataset(tfrecord_dir, res, batch_size, epochs)
+
+    for real_images in dataset.take(4):
+        # real_images: [batch_size, 3, res, res] (-1.0 ~ 1.0) float32
+        print(real_images.shape)
+
+        images = real_images.numpy()
+        images = np.transpose(images, axes=(0, 2, 3, 1))
+        images = (images + 1.0) * 127.5
+        images = images.astype(np.uint8)
+        image = Image.fromarray(images[0])
+        image.show()
+    return
+
+
+if __name__ == '__main__':
+    main()

docker-compose.yml

@@ -0,0 +1,40 @@
+version: '2.3'
+services:
+  training:
+    image: 172.20.40.190:5000/tf-2.0.0-gpu-py3:tfds
+    runtime: nvidia
+    environment:
+      - NVIDIA_VISIBLE_DEVICES=0,1
+    user: "2004:2004"
+    container_name: "moono_gpu_01"
+    volumes:
+      - "/mnt/vision-nas/moono/git-repos/stylegan2-tf-2.x:/work-dir"
+      - "/mnt/vision-nas/moono/trained_models/stylegan2-tf-2.x:/model-dir"
+      - "/mnt/vision-nas/data-sets/stylegan/ffhq-dataset/tfrecords/ffhq:/tfrecord-dir:ro"
+    working_dir: "/work-dir"
+    command: ["python", "-u", "train.py",
+              "--allow_memory_growth", "false",
+              "--debug_split_gpu", "false",
+              "--use_tf_function", "true",
+              "--use_custom_cuda", "true",
+              "--model_base_dir", "/model-dir/cuda",
+              "--tfrecord_dir", "/tfrecord-dir",
+              "--train_res", "256",
+              "--shuffle_buffer_size", "1000",
+              "--batch_size_per_replica", "16"
+    ]
+  watching:
+    image: 172.20.40.190:5000/tf-2.0.0-gpu-py3:tfds
+    runtime: nvidia
+    environment:
+      - NVIDIA_VISIBLE_DEVICES=1
+    user: "2004:2004"
+    container_name: "moono_tensorboard"
+    depends_on:
+      - training
+    volumes:
+      - "/mnt/vision-nas/moono/trained_models/stylegan2-tf-2.x:/model-dir"
+    ports:
+      - "36006:6006"
+    working_dir: "/model-dir"
+    command: ["tensorboard", "--logdir", "./", "--host", "0.0.0.0", "--port", "6006"]

docker/Dockerfile

@@ -0,0 +1,31 @@
+FROM tensorflow/tensorflow:2.3.1-gpu
+# FROM tensorflow/tensorflow:2.4.0rc0-gpu
+MAINTAINER moono.song "toilety@gmail.com"
+RUN apt-get update -y
+RUN apt-get install -y build-essential
+RUN pip install pillow
+
+# add sudoer
+ARG USERNAME=anonymous
+ARG USER_UID=1000
+ARG USER_GID=$USER_UID
+
+# Create the user
+RUN groupadd --gid $USER_GID $USERNAME \
+    && useradd --uid $USER_UID --gid $USER_GID -m $USERNAME \
+    #
+    # [Optional] Add sudo support. Omit if you don't need to install software after connecting.
+    && apt-get update \
+    && apt-get install -y sudo \
+    && echo $USERNAME ALL=\(root\) NOPASSWD:ALL > /etc/sudoers.d/$USERNAME \
+    && chmod 0440 /etc/sudoers.d/$USERNAME
+
+# ********************************************************
+# * Anything else you want to do like clean up goes here *
+# ********************************************************
+
+# [Optional] Set the default user. Omit if you want to keep the default as root.
+USER $USERNAME
+
+# add port for multiworker training
+EXPOSE 12345

exports/.gitignore

@@ -0,0 +1,5 @@
+# ignore all
+*
+
+# except this file
+!.gitignore

how_to_extract_official_weights.md

@@ -0,0 +1,85 @@
+# Steps to extract official weights
+
+### Prerequisite
+* install docker >= 19.03
+* install nvidia-docker native
+
+## Step 1. Working on official repository
+* Clone official repo
+```bash
+~$ git clone https://github.com/NVlabs/stylegan2.git
+~$ cd stylegan2
+```
+* Download official weights (*.pkl)
+* Create directory
+```bash
+~$ mkdir official_pretrained
+```
+* Place official weight (*.pkl) to `<PATH-TO-OFFICIAL-REPO>/official_pretrained`
+* Create python script called `export_weights.py` and add following
+```python
+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
+```bash
+~$ docker build -f Dockerfile -t stylegan2:official .
+``` 
+* Get inside docker container and run script
+```bash
+~$ 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
+```bash
+~$ exit
+```
+
+## Step 2. Working on this repository(stylegan2-tf-2.x)
+* Move to `stylegan2-tf-2.x` directory
+```bash
+~$ 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
+```bash
+~$ docker build -f Dockerfile -t stylegan2:tf.2.x .
+```
+* Get inside docker container and run script
+```bash
+~$ 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` 

inference.py

@@ -0,0 +1,65 @@
+import tensorflow as tf
+
+from PIL import Image
+from load_models import load_generator
+from tf_utils import allow_memory_growth
+
+
+def inference(ckpt_dir, use_custom_cuda, res, out_fn=None):
+    # create generator
+    resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
+    featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
+    filter_index = resolutions.index(res)
+    g_params = {
+        'z_dim': 512,
+        'w_dim': 512,
+        'labels_dim': 0,
+        'n_mapping': 8,
+        'resolutions': resolutions[:filter_index + 1],
+        'featuremaps': featuremaps[:filter_index + 1],
+    }
+    generator = load_generator(g_params, is_g_clone=True, ckpt_dir=ckpt_dir, custom_cuda=use_custom_cuda)
+
+    # generate image
+    fake_images = generator([tf.random.normal(shape=[1, g_params['z_dim']]),
+                             tf.random.normal(shape=[1, g_params['labels_dim']])],
+                            training=False, truncation_psi=0.5)
+    fake_images = (tf.clip_by_value(fake_images, -1.0, 1.0) + 1.0) * 127.5
+    fake_images = tf.transpose(fake_images, perm=[0, 2, 3, 1])
+    fake_images = tf.cast(fake_images, tf.uint8)
+    fake_image = fake_images[0].numpy()
+
+    image = Image.fromarray(fake_image)
+    image = image.convert('RGB')
+    image.show()
+    if out_fn is not None:
+        image.save(out_fn)
+    return
+
+
+def main():
+    allow_memory_growth()
+
+    checkpoints = [
+        {
+            'res': 512,
+            'ckpt_dir': '/workspace/Checkpoints/checkpoint_32100/stylegan2-test-2-512x512',
+            'use_custom_cuda': False,
+            'out_fn': 'out_0.png',
+        },
+    ]
+
+    for run_item in checkpoints:
+        res = run_item['res']
+        ckpt_dir = run_item['ckpt_dir']
+        use_custom_cuda = run_item['use_custom_cuda']
+        out_fn = run_item['out_fn']
+        message = f'{res}x{res} with custom cuda' if use_custom_cuda else f'{res}x{res} without custom cuda'
+
+        print(message)
+        inference(ckpt_dir, use_custom_cuda, res, out_fn)
+    return
+
+
+if __name__ == '__main__':
+    main()

inference_custom.py

@@ -0,0 +1,95 @@
+import tensorflow as tf
+
+import os
+import sys
+import cv2
+import numpy as np
+from PIL import Image
+from load_models import load_generator
+from tf_utils import allow_memory_growth
+
+
+def m2uv_array(m2uv):
+    result = np.zeros((14475, 2), dtype=np.int32)
+    for k, v in m2uv.items():
+        if len(v) == 2:
+            result[k] = list(v)
+        elif len(v) == 1:
+            x = v.pop()
+            result[k] = [x, x]
+        else:
+            assert False
+
+    return result
+
+
+def uv_to_mesh(uv):
+    uv_res = 256
+    m2uv = np.load('MeshUtils/m2uv.npy', allow_pickle=True).item()
+    m2uv = m2uv_array(m2uv)
+    uv_map = np.load('Meshutils/stretched_uv_map_3.npy')
+    uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
+
+    mesh = np.zeros((14475, 3))
+    uv_pixels = uv_map_discreet[m2uv]
+    # counter = np.ones((14475))
+    # counter[m2uv[:, 0] == m2uv[:, 1]] = 2
+    mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
+    return mesh
+
+
+def inference(ckpt_dir, use_custom_cuda, res, folder=None, samples=1000):
+    # create generator
+    resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
+    featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
+    filter_index = resolutions.index(res)
+    g_params = {
+        'z_dim': 512,
+        'w_dim': 512,
+        'labels_dim': 0,
+        'n_mapping': 8,
+        'resolutions': resolutions[:filter_index + 1],
+        'featuremaps': featuremaps[:filter_index + 1],
+    }
+    generator = load_generator(g_params, is_g_clone=True, ckpt_dir=ckpt_dir, custom_cuda=use_custom_cuda)
+
+    for i in range(samples):
+        print(i)
+        # generate image
+        fake_images = generator([tf.random.normal(shape=[1, g_params['z_dim']]),
+                                tf.random.normal(shape=[1, g_params['labels_dim']])],
+                                training=False, truncation_psi=0.5)
+
+        np.savez_compressed(os.path.join(folder, str(i+100) + '.npz'), uv=fake_images)
+        img = np.transpose(np.round((fake_images + 1) * 128)[0], [1,2,0])
+        cv2.imwrite(os.path.join(folder, str(i+100) + '.jpg'), img)
+
+    return
+
+
+def main():
+    allow_memory_growth()
+
+    checkpoints = [
+        {
+            'res': 512,
+            'ckpt_dir': '/workspace/Checkpoints/checkpoint_52000/stylegan2-test-2-512x512',
+            'use_custom_cuda': False,
+            'folder': '/workspace/Checkpoints/checkpoint_52000/output',
+        },
+    ]
+
+    for run_item in checkpoints:
+        res = run_item['res']
+        ckpt_dir = run_item['ckpt_dir']
+        use_custom_cuda = run_item['use_custom_cuda']
+        folder = run_item['folder']
+        message = f'{res}x{res} with custom cuda' if use_custom_cuda else f'{res}x{res} without custom cuda'
+
+        print(message)
+        inference(ckpt_dir, use_custom_cuda, res, folder)
+    return
+
+
+if __name__ == '__main__':
+    main()

inference_from_official_weights.py

@@ -0,0 +1,62 @@
+import os
+import numpy as np
+import tensorflow as tf
+
+from PIL import Image
+from stylegan2.utils import postprocess_images
+from load_models import load_generator
+from copy_official_weights import convert_official_weights_together
+
+
+def test_generator(ckpt_dir, use_custom_cuda, out_fn):
+    g_clone = load_generator(g_params=None, is_g_clone=True, ckpt_dir=ckpt_dir, custom_cuda=use_custom_cuda)
+
+    # test
+    seed = 6600
+    rnd = np.random.RandomState(seed)
+    latents = rnd.randn(1, g_clone.z_dim)
+    labels = rnd.randn(1, g_clone.labels_dim)
+    latents = latents.astype(np.float32)
+    labels = labels.astype(np.float32)
+    image_out = g_clone([latents, labels], training=False, truncation_psi=0.5)
+    image_out = postprocess_images(image_out)
+    image_out = image_out.numpy()
+
+    out_fn = f'seed{seed}-{out_fn}'
+    Image.fromarray(image_out[0], 'RGB').save(out_fn)
+    return
+
+
+def main():
+    from tf_utils import allow_memory_growth
+
+    allow_memory_growth()
+
+    # common variables
+    ckpt_dir_base = './official-converted'
+
+    # saving phase
+    for use_custom_cuda in [True, False]:
+        ckpt_dir = os.path.join(ckpt_dir_base, 'cuda') if use_custom_cuda else os.path.join(ckpt_dir_base, 'ref')
+        convert_official_weights_together(ckpt_dir, use_custom_cuda)
+
+    # inference phase
+    ckpt_dir_cuda = os.path.join(ckpt_dir_base, 'cuda')
+    ckpt_dir_ref = os.path.join(ckpt_dir_base, 'ref')
+
+    # 1. inference cuda saved weight from cuda model
+    test_generator(ckpt_dir_cuda, use_custom_cuda=True, out_fn='from-cuda-to-cuda.png')
+
+    # 2. inference cuda saved weight from ref model
+    test_generator(ckpt_dir_cuda, use_custom_cuda=False, out_fn='from-cuda-to-ref.png')
+
+    # 3. inference ref saved weight from ref model
+    test_generator(ckpt_dir_ref, use_custom_cuda=False, out_fn='from-ref-to-ref.png')
+
+    # 4. inference ref saved weight from cuda model
+    test_generator(ckpt_dir_ref, use_custom_cuda=True, out_fn='from-ref-to-cuda.png')
+    return
+
+
+if __name__ == '__main__':
+    main()

load_models.py

@@ -0,0 +1,66 @@
+import tensorflow as tf
+
+
+def load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=True):
+    if custom_cuda:
+        from stylegan2.generator import Generator
+    else:
+        from stylegan2_ref.generator import Generator
+
+    if g_params is None:
+        g_params = {
+            'z_dim': 512,
+            'w_dim': 512,
+            'labels_dim': 0,
+            'n_mapping': 8,
+            'resolutions': [4, 8, 16, 32, 64, 128, 256, 512, 1024],
+            'featuremaps': [512, 512, 512, 512, 512, 256, 128, 64, 32],
+        }
+
+    test_latent = tf.ones((1, g_params['z_dim']), dtype=tf.float32)
+    test_labels = tf.ones((1, g_params['labels_dim']), dtype=tf.float32)
+
+    # build generator model
+    generator = Generator(g_params)
+    _ = generator([test_latent, test_labels])
+
+    if ckpt_dir is not None:
+        if is_g_clone:
+            ckpt = tf.train.Checkpoint(g_clone=generator)
+        else:
+            ckpt = tf.train.Checkpoint(generator=generator)
+        manager = tf.train.CheckpointManager(ckpt, ckpt_dir, max_to_keep=1)
+        ckpt.restore(manager.latest_checkpoint).expect_partial()
+        if manager.latest_checkpoint:
+            print(f'Generator restored from {manager.latest_checkpoint}')
+    return generator
+
+
+def load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=True):
+    if custom_cuda:
+        from stylegan2.discriminator import Discriminator
+    else:
+        from stylegan2_ref.discriminator import Discriminator
+
+    if d_params is None:
+        d_params = {
+            'labels_dim': 0,
+            'resolutions': [4, 8, 16, 32, 64, 128, 256, 512, 1024],
+            'featuremaps': [512, 512, 512, 512, 512, 256, 128, 64, 32],
+        }
+
+    res = d_params['resolutions'][-1]
+    test_images = tf.ones((1, 3, res, int(res/2)), dtype=tf.float32)
+    test_labels = tf.ones((1, d_params['labels_dim']), dtype=tf.float32)
+
+    # build discriminator model
+    discriminator = Discriminator(d_params)
+    _ = discriminator([test_images, test_labels])
+
+    if ckpt_dir is not None:
+        ckpt = tf.train.Checkpoint(discriminator=discriminator)
+        manager = tf.train.CheckpointManager(ckpt, ckpt_dir, max_to_keep=1)
+        ckpt.restore(manager.latest_checkpoint).expect_partial()
+        if manager.latest_checkpoint:
+            print('Discriminator restored from {}'.format(manager.latest_checkpoint))
+    return discriminator

losses.py

@@ -0,0 +1,96 @@
+import tensorflow as tf
+
+
+def d_logistic(real_images, generator, discriminator, z_dim, labels):
+    batch_size = tf.shape(real_images)[0]
+    z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
+
+    # forward pass
+    fake_images = generator([z, labels], training=True)
+    real_scores, real_cond = discriminator([real_images, labels], training=True)
+    fake_scores, fake_cond = discriminator([fake_images, labels], training=True)
+
+    # gan loss
+    d_loss = tf.math.softplus(fake_scores)
+    d_loss += tf.math.softplus(-real_scores)
+
+    cond_loss = tf.nn.l2_loss(labels - real_cond) + tf.nn.l2_loss(labels - fake_cond)
+
+    return d_loss, cond_loss
+
+
+def d_logistic_r1_reg(real_images, generator, discriminator, z_dim, labels):
+    batch_size = tf.shape(real_images)[0]
+    z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
+
+    # forward pass
+    fake_images = generator([z, labels], training=True)
+    real_scores, real_cond = discriminator([real_images, labels], training=True)
+    fake_scores, fake_cond = discriminator([fake_images, labels], training=True)
+
+    cond_loss = tf.nn.l2_loss(labels - real_cond) + tf.nn.l2_loss(labels - fake_cond)
+
+    # gan loss
+    d_loss = tf.math.softplus(fake_scores)
+    d_loss = d_loss + tf.math.softplus(-real_scores)
+
+    # gradient penalty
+    with tf.GradientTape() as r1_tape:
+        r1_tape.watch([real_images, labels])
+        real_scores, _ = discriminator([real_images, labels], training=True)
+        real_loss = tf.reduce_sum(real_scores)
+
+    real_grads = r1_tape.gradient(real_loss, real_images)
+    r1_penalty = tf.reduce_sum(tf.math.square(real_grads), axis=[1, 2, 3])
+    r1_penalty = tf.expand_dims(r1_penalty, axis=1)
+    return d_loss, cond_loss, r1_penalty
+
+
+def g_logistic_non_saturating(real_images, generator, discriminator, z_dim, labels):
+    batch_size = tf.shape(real_images)[0]
+    z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
+
+    # forward pass
+    fake_images = generator([z, labels], training=True)
+    fake_scores, _ = discriminator([fake_images, labels], training=True)
+    # TODO: do something with cond output?
+
+    # gan loss
+    g_loss = tf.math.softplus(-fake_scores)
+    return g_loss
+
+
+def g_logistic_ns_pathreg(real_images, generator, discriminator, z_dim,
+                          pl_mean, pl_minibatch_shrink, pl_denorm, pl_decay,
+                          labels):
+    batch_size = tf.shape(real_images)[0]
+    z = tf.random.normal(shape=[batch_size, z_dim], dtype=tf.float32)
+
+    pl_minibatch = tf.maximum(1, tf.math.floordiv(batch_size, pl_minibatch_shrink))
+    pl_z = tf.random.normal(shape=[pl_minibatch, z_dim], dtype=tf.float32)
+    pl_labels = labels[:pl_minibatch]
+
+    # forward pass
+    fake_images = generator([z, labels], ret_w_broadcasted=False, training=True)
+    fake_scores, _ = discriminator([fake_images, labels], training=True)
+    # TODO: do something with cond output?
+    g_loss = tf.math.softplus(-fake_scores)
+
+    # Evaluate the regularization term using a smaller minibatch to conserve memory.
+    with tf.GradientTape() as pl_tape:
+        pl_tape.watch([pl_z, pl_labels])
+        pl_fake_images, pl_w_broadcasted = generator([pl_z, pl_labels], ret_w_broadcasted=True, training=True)
+
+        pl_noise = tf.random.normal(tf.shape(pl_fake_images)) * pl_denorm
+        pl_noise_applied = tf.reduce_sum(pl_fake_images * pl_noise)
+
+    pl_grads = pl_tape.gradient(pl_noise_applied, pl_w_broadcasted)
+    pl_lengths = tf.math.sqrt(tf.reduce_mean(tf.reduce_sum(tf.math.square(pl_grads), axis=2), axis=1))
+
+    # Track exponential moving average of |J*y|.
+    pl_mean_val = pl_mean + pl_decay * (tf.reduce_mean(pl_lengths) - pl_mean)
+    pl_mean.assign(pl_mean_val)
+
+    # Calculate (|J*y|-a)^2.
+    pl_penalty = tf.square(pl_lengths - pl_mean)
+    return g_loss, pl_penalty

models/.gitignore

@@ -0,0 +1,5 @@
+# ignore all
+*
+
+# except this file
+!.gitignore

official-converted/.gitignore

@@ -0,0 +1,5 @@
+# ignore all
+*
+
+# except this file
+!.gitignore

official-pretrained/.gitignore

@@ -0,0 +1,5 @@
+# ignore all
+*
+
+# except this file
+!.gitignore

reconstruct.bat

@@ -0,0 +1 @@
+"D:\Work\Code\env\Scripts\python.exe" D:\Work\Code\stylegan2-tf-2.x\reconstruct.py %*

reconstruct.py

@@ -0,0 +1,71 @@
+import os
+import sys
+import cv2
+import numpy as np
+import trimesh as tm
+import matplotlib.pyplot as plt
+
+def m2uv_array(m2uv):
+    result = np.zeros((14475, 2), dtype=np.int32)
+    for k, v in m2uv.items():
+        if len(v) == 2:
+            result[k] = list(v)
+        elif len(v) == 1:
+            x = v.pop()
+            result[k] = [x, x]
+        else:
+            assert False
+
+    return result
+
+
+def uv_to_mesh(uv):
+    uv_res = 256
+    # m2uv = np.load('/workspace/Code/featurelines/MeshUtils/m2uv.npy', allow_pickle=True).item()
+    m2uv = np.load(r'D:\Work\Code\featurelines\MeshUtils\m2uv.npy', allow_pickle=True).item()
+    m2uv = m2uv_array(m2uv)
+    # uv_map = np.load('/workspace/Code/featurelines/Meshutils/stretched_uv_map_3.npy')
+    uv_map = np.load(r'D:\Work\Code\featurelines\MeshUtils\stretched_uv_map_3.npy')
+    uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
+
+    mesh = np.zeros((14475, 3))
+    uv_pixels = uv_map_discreet[m2uv]
+    # counter = np.ones((14475))
+    # counter[m2uv[:, 0] == m2uv[:, 1]] = 2
+    mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
+    return mesh
+
+
+def plot(uv, mesh):
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(111, projection='3d')
+    ax.set_xlim3d(-1, 1)
+    ax.set_ylim3d(-1, 1)
+    ax.set_zlim3d(-1, 1)
+    ax.set_xlabel("X")
+    ax.set_ylabel("Y")
+    ax.set_zlabel("Z")
+    ax.scatter(mesh[:, 0], mesh[:, 1], mesh[:, 2], s=5)
+    # plt.show()
+
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(111)
+    ax.imshow(np.round( (uv - uv.min()) / (uv.max() - uv.min()) * 255).astype(np.int32))
+    plt.show()
+
+    F = np.load(r'D:\Work\Code\featurelines\FeatureLines\superSMPL\smpl_faces.npy')
+    tm.Trimesh(mesh, F).show()
+
+
+
+if __name__ == '__main__':
+    file = sys.argv[1]
+
+    uv = np.load(file)
+    uv = np.transpose(uv['uv'][0], [1,2,0])
+    mesh = uv_to_mesh(uv)
+
+    plot(uv, mesh)
+
+    pass
+

recontruct_mesh.py

@@ -0,0 +1,77 @@
+import os
+import sys
+import cv2
+import numpy as np
+import trimesh as tm
+import matplotlib.pyplot as plt
+
+def m2uv_array(m2uv):
+    result = np.zeros((14475, 2), dtype=np.int32)
+    for k, v in m2uv.items():
+        if len(v) == 2:
+            result[k] = list(v)
+        elif len(v) == 1:
+            x = v.pop()
+            result[k] = [x, x]
+        else:
+            assert False
+
+    return result
+
+
+def uv_to_mesh(uv):
+    uv_res = 256
+    # m2uv = np.load('/workspace/Code/featurelines/MeshUtils/m2uv.npy', allow_pickle=True).item()
+    m2uv = np.load(r'D:\Work\Code\featurelines\MeshUtils\m2uv.npy', allow_pickle=True).item()
+    m2uv = m2uv_array(m2uv)
+    # uv_map = np.load('/workspace/Code/featurelines/Meshutils/stretched_uv_map_3.npy')
+    uv_map = np.load(r'D:\Work\Code\featurelines\MeshUtils\stretched_uv_map_3.npy')
+    uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
+
+    mesh = np.zeros((14475, 3))
+    uv_pixels = uv_map_discreet[m2uv]
+    # counter = np.ones((14475))
+    # counter[m2uv[:, 0] == m2uv[:, 1]] = 2
+    mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
+    return mesh
+
+
+def plot(uv, mesh):
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(111, projection='3d')
+    ax.set_xlim3d(-1, 1)
+    ax.set_ylim3d(-1, 1)
+    ax.set_zlim3d(-1, 1)
+    ax.set_xlabel("X")
+    ax.set_ylabel("Y")
+    ax.set_zlabel("Z")
+    ax.scatter(mesh[:, 0], mesh[:, 1], mesh[:, 2], s=5)
+    # plt.show()
+
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(111)
+    ax.imshow(np.round( (uv - uv.min()) / (uv.max() - uv.min()) * 255).astype(np.int32))
+    plt.show()
+
+
+
+if __name__ == '__main__':
+    # folder =  '/workspace/Checkpoints/checkpoint_32100/output'
+    folder = r'D:\Work\Checkpoints\checkpoint_52000\output'
+
+    files = os.listdir(folder)
+
+    for file in files:
+        if not file.endswith('.npz'):
+            continue
+
+        uv = np.load(os.path.join(folder, file))
+        uv = np.transpose(uv['uv'][0], [1,2,0])
+        # uv[abs(uv) < 0.01] = 0
+        # print(np.sum(abs(uv) < 0.01))
+        mesh = uv_to_mesh(uv)
+
+        plot(uv, mesh)
+
+        pass
+

stylegan2/discriminator.py

@@ -0,0 +1,162 @@
+import tensorflow as tf
+
+from stylegan2.layers.dense import Dense
+from stylegan2.layers.conv import Conv2D
+from stylegan2.layers.bias_act import BiasAct
+from stylegan2.layers.from_rgb import FromRGB
+from stylegan2.layers.mini_batch_std import MinibatchStd
+
+
+class DiscriminatorBlock(tf.keras.layers.Layer):
+    def __init__(self, n_f0, n_f1, res, **kwargs):
+        super(DiscriminatorBlock, self).__init__(**kwargs)
+        self.gain = 1.0
+        self.lrmul = 1.0
+        self.n_f0 = n_f0
+        self.n_f1 = n_f1
+        self.res = res
+        self.resnet_scale = 1. / tf.sqrt(2.)
+
+        # conv_0
+        self.conv_0 = Conv2D(in_res=res, in_fmaps=self.n_f0, fmaps=self.n_f0, kernel=3, up=False, down=False,
+                             resample_kernel=None, gain=self.gain, lrmul=self.lrmul, name='conv_0')
+        self.apply_bias_act_0 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_0')
+
+        # conv_1 down
+        self.conv_1 = Conv2D(in_res=res, in_fmaps=self.n_f0, fmaps=self.n_f1, kernel=3, up=False, down=True,
+                             resample_kernel=[1, 3, 3, 1], gain=self.gain, lrmul=self.lrmul, name='conv_1')
+        self.apply_bias_act_1 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_1')
+
+        # resnet skip
+        self.conv_skip = Conv2D(in_res=res, in_fmaps=self.n_f0, fmaps=self.n_f1, kernel=1, up=False, down=True,
+                                resample_kernel=[1, 3, 3, 1], gain=self.gain, lrmul=self.lrmul, name='skip')
+
+    def call(self, inputs, training=None, mask=None):
+        x = inputs
+        residual = x
+
+        # conv0
+        x = self.conv_0(x)
+        x = self.apply_bias_act_0(x)
+
+        # conv1 down
+        x = self.conv_1(x)
+        x = self.apply_bias_act_1(x)
+
+        # resnet skip
+        residual = self.conv_skip(residual)
+        x = (x + residual) * self.resnet_scale
+        return x
+
+    def get_config(self):
+        config = super(DiscriminatorBlock, self).get_config()
+        config.update({
+            'n_f0': self.n_f0,
+            'n_f1': self.n_f1,
+            'gain': self.gain,
+            'lrmul': self.lrmul,
+            'res': self.res,
+            'resnet_scale': self.resnet_scale,
+        })
+        return config
+
+
+class DiscriminatorLastBlock(tf.keras.layers.Layer):
+    def __init__(self, n_f0, n_f1, res, **kwargs):
+        super(DiscriminatorLastBlock, self).__init__(**kwargs)
+        self.gain = 1.0
+        self.lrmul = 1.0
+        self.n_f0 = n_f0
+        self.n_f1 = n_f1
+        self.res = res
+
+        self.minibatch_std = MinibatchStd(group_size=4, num_new_features=1, name='minibatchstd')
+
+        # conv_0
+        self.conv_0 = Conv2D(in_res=res, in_fmaps=self.n_f0 + 1, fmaps=self.n_f0, kernel=3, up=False, down=False,
+                             resample_kernel=None, gain=self.gain, lrmul=self.lrmul, name='conv_0')
+        self.apply_bias_act_0 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_0')
+
+        # dense_1
+        self.dense_1 = Dense(self.n_f1, gain=self.gain, lrmul=self.lrmul, name='dense_1')
+        self.apply_bias_act_1 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_1')
+
+    def call(self, x, training=None, mask=None):
+        x = self.minibatch_std(x)
+
+        # conv_0
+        x = self.conv_0(x)
+        x = self.apply_bias_act_0(x)
+
+        # dense_1
+        x = self.dense_1(x)
+        x = self.apply_bias_act_1(x)
+        return x
+
+    def get_config(self):
+        config = super(DiscriminatorLastBlock, self).get_config()
+        config.update({
+            'n_f0': self.n_f0,
+            'n_f1': self.n_f1,
+            'gain': self.gain,
+            'lrmul': self.lrmul,
+            'res': self.res,
+        })
+        return config
+
+
+class Discriminator(tf.keras.Model):
+    def __init__(self, d_params, **kwargs):
+        super(Discriminator, self).__init__(**kwargs)
+        # discriminator's (resolutions and featuremaps) are reversed against generator's
+        self.labels_dim = d_params['labels_dim']
+        self.r_resolutions = d_params['resolutions'][::-1]
+        self.r_featuremaps = d_params['featuremaps'][::-1]
+
+        # stack discriminator blocks
+        res0, n_f0 = self.r_resolutions[0], self.r_featuremaps[0]
+        self.initial_fromrgb = FromRGB(fmaps=n_f0, res=res0, name='{:d}x{:d}/FromRGB'.format(res0, res0))
+        self.blocks = list()
+        for index, (res0, n_f0) in enumerate(zip(self.r_resolutions[:-1], self.r_featuremaps[:-1])):
+            n_f1 = self.r_featuremaps[index + 1]
+            self.blocks.append(DiscriminatorBlock(n_f0=n_f0, n_f1=n_f1, res=res0, name='{:d}x{:d}'.format(res0, res0)))
+
+        # set last discriminator block
+        res = self.r_resolutions[-1]
+        n_f0, n_f1 = self.r_featuremaps[-2], self.r_featuremaps[-1]
+        self.last_block = DiscriminatorLastBlock(n_f0, n_f1, res, name='{:d}x{:d}'.format(res, res))
+
+        # set last dense layer
+        self.last_dense = Dense(max(self.labels_dim, 1), gain=1.0, lrmul=1.0, name='last_dense')
+        self.last_bias = BiasAct(lrmul=1.0, act='linear', name='last_bias')
+
+        # conditional output
+        self.last_block_condition = DiscriminatorLastBlock(n_f0, n_f1, res, name='{:d}x{:d}'.format(res, res))
+        self.condition_dense = Dense(self.labels_dim, gain=1.0, lrmul=1.0, name='last_dense')
+        self.condition_bias = BiasAct(lrmul=1.0, act='linear', name='last_bias')
+
+
+    def call(self, inputs, training=None, mask=None):
+        images, _ = inputs
+
+        x = self.initial_fromrgb(images)
+        for block in self.blocks:
+            x = block(x)
+
+        features = x
+
+        # score output
+        x = self.last_block(features)
+        x = self.last_dense(x)
+        x = self.last_bias(x)
+
+        # condition output
+        cond = self.last_block_condition(features)
+        cond = self.condition_dense(cond)
+        cond = self.condition_bias(cond)
+
+        scores_out = x
+        return scores_out, cond
+
+    def compute_output_shape(self, input_shape):
+        return input_shape[0][0], max(self.labels_dim, 1)

stylegan2/generator.py

@@ -0,0 +1,112 @@
+import tensorflow as tf
+
+from stylegan2.utils import lerp
+from stylegan2.layers.mapping_block import Mapping
+from stylegan2.layers.synthesis_block import Synthesis
+
+
+class Generator(tf.keras.Model):
+    def __init__(self, g_params, **kwargs):
+        super(Generator, self).__init__(**kwargs)
+
+        self.z_dim = g_params['z_dim']
+        self.w_dim = g_params['w_dim']
+        self.labels_dim = g_params['labels_dim']
+        self.n_mapping = g_params['n_mapping']
+        self.resolutions = g_params['resolutions']
+        self.featuremaps = g_params['featuremaps']
+        self.w_ema_decay = 0.995
+        self.style_mixing_prob = 0.9
+
+        self.n_broadcast = len(self.resolutions) * 2
+        self.mixing_layer_indices = tf.range(self.n_broadcast, dtype=tf.int32)[tf.newaxis, :, tf.newaxis]
+
+        self.g_mapping = Mapping(self.w_dim, self.labels_dim, self.n_mapping, name='g_mapping')
+        self.broadcast = tf.keras.layers.Lambda(lambda x: tf.tile(x[:, tf.newaxis], [1, self.n_broadcast, 1]))
+        self.synthesis = Synthesis(self.resolutions, self.featuremaps, name='g_synthesis')
+
+    def build(self, input_shape):
+        # w_avg
+        self.w_avg = tf.Variable(tf.zeros(shape=[self.w_dim], dtype=tf.dtypes.float32), name='w_avg', trainable=False,
+                                 synchronization=tf.VariableSynchronization.ON_READ,
+                                 aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA)
+
+    @tf.function
+    def set_as_moving_average_of(self, src_net):
+        beta, beta_nontrainable = 0.99, 0.0
+
+        for cw, sw in zip(self.weights, src_net.weights):
+            assert sw.shape == cw.shape
+            # print('{} <=> {}'.format(cw.name, sw.name))
+
+            if 'w_avg' in cw.name:
+                cw.assign(lerp(sw, cw, beta_nontrainable))
+            else:
+                cw.assign(lerp(sw, cw, beta))
+        return
+
+    def update_moving_average_of_w(self, w_broadcasted):
+        # compute average of current w
+        batch_avg = tf.reduce_mean(w_broadcasted[:, 0], axis=0)
+
+        # compute moving average of w and update(assign) w_avg
+        self.w_avg.assign(lerp(batch_avg, self.w_avg, self.w_ema_decay))
+        return
+
+    def style_mixing_regularization(self, latents1, labels, w_broadcasted1):
+        # get another w and broadcast it
+        latents2 = tf.random.normal(shape=tf.shape(latents1), dtype=tf.dtypes.float32)
+        dlatents2 = self.g_mapping([latents2, labels])
+        w_broadcasted2 = self.broadcast(dlatents2)
+
+        # find mixing limit index
+        # mixing_cutoff_index = tf.cond(
+        #     pred=tf.less(tf.random.uniform([], 0.0, 1.0), self.style_mixing_prob),
+        #     true_fn=lambda: tf.random.uniform([], 1, self.n_broadcast, dtype=tf.dtypes.int32),
+        #     false_fn=lambda: tf.constant(self.n_broadcast, dtype=tf.dtypes.int32))
+        if tf.random.uniform([], 0.0, 1.0) < self.style_mixing_prob:
+            mixing_cutoff_index = tf.random.uniform([], 1, self.n_broadcast, dtype=tf.dtypes.int32)
+        else:
+            mixing_cutoff_index = tf.constant(self.n_broadcast, dtype=tf.dtypes.int32)
+
+        # mix it
+        mixed_w_broadcasted = tf.where(
+            condition=tf.broadcast_to(self.mixing_layer_indices < mixing_cutoff_index, tf.shape(w_broadcasted1)),
+            x=w_broadcasted1,
+            y=w_broadcasted2)
+        return mixed_w_broadcasted
+
+    def truncation_trick(self, w_broadcasted, truncation_psi, truncation_cutoff=None):
+        ones = tf.ones_like(self.mixing_layer_indices, dtype=tf.float32)
+        tpsi = ones * truncation_psi
+        if truncation_cutoff is None:
+            truncation_coefs = tpsi
+        else:
+            indices = tf.range(self.n_broadcast)
+            truncation_coefs = tf.where(condition=tf.less(indices, truncation_cutoff), x=tpsi, y=ones)
+
+        truncated_w_broadcasted = lerp(self.w_avg, w_broadcasted, truncation_coefs)
+        return truncated_w_broadcasted
+
+    def call(self, inputs, ret_w_broadcasted=False, truncation_psi=1.0, truncation_cutoff=None, training=None, mask=None):
+        latents, labels = inputs
+
+        dlatents = self.g_mapping([latents, labels])
+        w_broadcasted = self.broadcast(dlatents)
+
+        if training:
+            self.update_moving_average_of_w(w_broadcasted)
+            w_broadcasted = self.style_mixing_regularization(latents, labels, w_broadcasted)
+
+        if not training:
+            w_broadcasted = self.truncation_trick(w_broadcasted, truncation_psi, truncation_cutoff)
+
+        image_out = self.synthesis(w_broadcasted)
+
+        if ret_w_broadcasted:
+            return image_out, w_broadcasted
+        else:
+            return image_out
+
+    def compute_output_shape(self, input_shape):
+        return input_shape[0][0], 3, self.resolutions[-1], self.resolutions[-1]

stylegan2/layers/bias_act.py

@@ -0,0 +1,73 @@
+import numpy as np
+import tensorflow as tf
+
+
+# the custom cuda compiled fused_bias_act is not compatible with tf.GradientTape()
+# I think it's because of 2nd order gradient computation
+class BiasAct(tf.keras.layers.Layer):
+    def __init__(self, lrmul, act, **kwargs):
+        super(BiasAct, self).__init__(**kwargs)
+        assert act in ['linear', 'lrelu']
+        self.lrmul = lrmul
+
+        if act == 'linear':
+            self.act = tf.keras.layers.Lambda(lambda x: tf.identity(x))
+            self.gain = 1.0
+        else:
+            self.act = tf.keras.layers.LeakyReLU(alpha=0.2)
+            self.gain = np.sqrt(2)
+
+    def build(self, input_shape):
+        self.len2 = True if len(input_shape) == 2 else False
+        b_init = tf.zeros(shape=(input_shape[1],), dtype=tf.dtypes.float32)
+        self.b = tf.Variable(b_init, name='b', trainable=True)
+
+    def call(self, inputs, training=None, mask=None):
+        b = self.lrmul * self.b
+
+        if self.len2:
+            x = inputs + b
+        else:
+            x = inputs + tf.reshape(b, shape=[1, -1, 1, 1])
+        x = self.act(x)
+        x = self.gain * x
+        return x
+
+    def get_config(self):
+        config = super(BiasAct, self).get_config()
+        config.update({
+            'lrmul': self.lrmul,
+            'gain': self.gain,
+            'len2': self.len2,
+        })
+        return config
+
+
+# import tensorflow as tf
+#
+#
+# from stylegan2.layers.cuda.fused_bias_act import fused_bias_act
+#
+#
+# class BiasAct(tf.keras.layers.Layer):
+#     def __init__(self, lrmul, act, **kwargs):
+#         super(BiasAct, self).__init__(**kwargs)
+#         self.lrmul = lrmul
+#         self.act = act
+#
+#     def build(self, input_shape):
+#         b_init = tf.zeros(shape=(input_shape[1],), dtype=tf.dtypes.float32)
+#         self.b = tf.Variable(b_init, name='b', trainable=True)
+#
+#     def call(self, inputs, training=None, mask=None):
+#         b = self.lrmul * self.b
+#         x = fused_bias_act(inputs, b=b, act=self.act, alpha=None, gain=None)
+#         return x
+#
+#     def get_config(self):
+#         config = super(BiasAct, self).get_config()
+#         config.update({
+#             'lrmul': self.lrmul,
+#             'act': self.act,
+#         })
+#         return config

stylegan2/layers/commons.py

@@ -0,0 +1,20 @@
+import tensorflow as tf
+
+
+def compute_runtime_coef(weight_shape, gain, lrmul):
+    fan_in = tf.reduce_prod(weight_shape[:-1])  # [kernel, kernel, fmaps_in, fmaps_out] or [in, out]
+    fan_in = tf.cast(fan_in, dtype=tf.float32)
+    he_std = gain / tf.sqrt(fan_in)
+    init_std = 1.0 / lrmul
+    runtime_coef = he_std * lrmul
+    return init_std, runtime_coef
+
+
+def lerp(a, b, t):
+    out = a + (b - a) * t
+    return out
+
+
+def lerp_clip(a, b, t):
+    out = a + (b - a) * tf.clip_by_value(t, 0.0, 1.0)
+    return out

stylegan2/layers/conv.py

@@ -0,0 +1,59 @@
+import tensorflow as tf
+
+from stylegan2.layers.commons import compute_runtime_coef
+from stylegan2.layers.cuda.upfirdn_2d_v2 import upsample_conv_2d, conv_downsample_2d, compute_paddings
+
+
+class Conv2D(tf.keras.layers.Layer):
+    def __init__(self, in_res, in_fmaps, fmaps, kernel, up, down, resample_kernel, gain, lrmul, **kwargs):
+        super(Conv2D, self).__init__(**kwargs)
+        self.in_res = in_res
+        self.in_fmaps = in_fmaps
+        self.fmaps = fmaps
+        self.kernel = kernel
+        self.gain = gain
+        self.lrmul = lrmul
+        self.up = up
+        self.down = down
+        # self.resample_kernel = resample_kernel
+
+        self.k, self.pad0, self.pad1 = compute_paddings(resample_kernel, self.kernel, up, down, is_conv=True)
+
+    def build(self, input_shape):
+        weight_shape = [self.kernel, self.kernel, input_shape[1], self.fmaps]
+        init_std, self.runtime_coef = compute_runtime_coef(weight_shape, self.gain, self.lrmul)
+
+        # [kernel, kernel, fmaps_in, fmaps_out]
+        w_init = tf.random.normal(shape=weight_shape, mean=0.0, stddev=init_std)
+        self.w = tf.Variable(w_init, name='w', trainable=True)
+
+    def call(self, inputs, training=None, mask=None):
+        x = inputs
+        w = self.runtime_coef * self.w
+
+        # actual conv
+        if self.up:
+            x = upsample_conv_2d(x, self.in_res, w, self.kernel, self.kernel, self.pad0, self.pad1, self.k)
+        elif self.down:
+            x = conv_downsample_2d(x, self.in_res, w, self.kernel, self.kernel, self.pad0, self.pad1, self.k)
+        else:
+            x = tf.nn.conv2d(x, w, data_format='NCHW', strides=[1, 1, 1, 1], padding='SAME')
+        return x
+
+    def get_config(self):
+        config = super(Conv2D, self).get_config()
+        config.update({
+            'in_res': self.in_res,
+            'in_fmaps': self.in_fmaps,
+            'fmaps': self.fmaps,
+            'kernel': self.kernel,
+            'gain': self.gain,
+            'lrmul': self.lrmul,
+            'up': self.up,
+            'down': self.down,
+            'k': self.k,
+            'pad0': self.pad0,
+            'pad1': self.pad1,
+            'runtime_coef': self.runtime_coef,
+        })
+        return config

stylegan2/layers/cuda/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://nvlabs.github.io/stylegan2/license.html
+
+# empty

stylegan2/layers/cuda/custom_ops.py

@@ -0,0 +1,169 @@
+# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://nvlabs.github.io/stylegan2/license.html
+
+"""TensorFlow custom ops builder.
+"""
+
+import os
+import re
+import uuid
+import hashlib
+import tempfile
+import shutil
+import tensorflow as tf
+from tensorflow.python.client import device_lib # pylint: disable=no-name-in-module
+
+#----------------------------------------------------------------------------
+# Global options.
+
+cuda_cache_path = os.path.join(os.path.dirname(__file__), '_cudacache')
+cuda_cache_version_tag = 'v1'
+do_not_hash_included_headers = False # Speed up compilation by assuming that headers included by the CUDA code never change. Unsafe!
+verbose = True # Print status messages to stdout.
+
+compiler_bindir_search_path = [
+    'C:/Program Files (x86)/Microsoft Visual Studio/2017/Community/VC/Tools/MSVC/14.14.26428/bin/Hostx64/x64',
+    'C:/Program Files (x86)/Microsoft Visual Studio/2019/Community/VC/Tools/MSVC/14.23.28105/bin/Hostx64/x64',
+    'C:/Program Files (x86)/Microsoft Visual Studio 14.0/vc/bin',
+]
+
+#----------------------------------------------------------------------------
+# Internal helper funcs.
+
+def _find_compiler_bindir():
+    for compiler_path in compiler_bindir_search_path:
+        if os.path.isdir(compiler_path):
+            return compiler_path
+    return None
+
+def _get_compute_cap(device):
+    caps_str = device.physical_device_desc
+    m = re.search('compute capability: (\\d+).(\\d+)', caps_str)
+    major = m.group(1)
+    minor = m.group(2)
+    return (major, minor)
+
+def _get_cuda_gpu_arch_string():
+    gpus = [x for x in device_lib.list_local_devices() if x.device_type == 'GPU']
+    if len(gpus) == 0:
+        raise RuntimeError('No GPU devices found')
+    (major, minor) = _get_compute_cap(gpus[0])
+    return 'sm_%s%s' % (major, minor)
+
+def _run_cmd(cmd):
+    with os.popen(cmd) as pipe:
+        output = pipe.read()
+        status = pipe.close()
+    if status is not None:
+        raise RuntimeError('NVCC returned an error. See below for full command line and output log:\n\n%s\n\n%s' % (cmd, output))
+
+def _prepare_nvcc_cli(opts):
+    cmd = 'nvcc --std=c++14 -DNDEBUG ' + opts.strip()
+    cmd += ' --disable-warnings'
+    cmd += ' --include-path "%s"' % tf.sysconfig.get_include()
+    cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'protobuf_archive', 'src')
+    cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'com_google_absl')
+    cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'eigen_archive')
+
+    compiler_bindir = _find_compiler_bindir()
+    if compiler_bindir is None:
+        # Require that _find_compiler_bindir succeeds on Windows.  Allow
+        # nvcc to use whatever is the default on Linux.
+        if os.name == 'nt':
+            raise RuntimeError('Could not find MSVC/GCC/CLANG installation on this computer. Check compiler_bindir_search_path list in "%s".' % __file__)
+    else:
+        cmd += ' --compiler-bindir "%s"' % compiler_bindir
+    cmd += ' 2>&1'
+    return cmd
+
+#----------------------------------------------------------------------------
+# Main entry point.
+
+_plugin_cache = dict()
+
+def get_plugin(cuda_file):
+    cuda_file_base = os.path.basename(cuda_file)
+    cuda_file_name, cuda_file_ext = os.path.splitext(cuda_file_base)
+
+    # Already in cache?
+    if cuda_file in _plugin_cache:
+        return _plugin_cache[cuda_file]
+
+    # Setup plugin.
+    if verbose:
+        print('Setting up TensorFlow plugin "%s": ' % cuda_file_base, end='', flush=True)
+    try:
+        # Hash CUDA source.
+        md5 = hashlib.md5()
+        with open(cuda_file, 'rb') as f:
+            md5.update(f.read())
+        md5.update(b'\n')
+
+        # Hash headers included by the CUDA code by running it through the preprocessor.
+        if not do_not_hash_included_headers:
+            if verbose:
+                print('Preprocessing... ', end='', flush=True)
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                tmp_file = os.path.join(tmp_dir, cuda_file_name + '_tmp' + cuda_file_ext)
+                _run_cmd(_prepare_nvcc_cli('"%s" --preprocess -o "%s" --keep --keep-dir "%s"' % (cuda_file, tmp_file, tmp_dir)))
+                with open(tmp_file, 'rb') as f:
+                    bad_file_str = ('"' + cuda_file.replace('\\', '/') + '"').encode('utf-8') # __FILE__ in error check macros
+                    good_file_str = ('"' + cuda_file_base + '"').encode('utf-8')
+                    for ln in f:
+                        if not ln.startswith(b'# ') and not ln.startswith(b'#line '): # ignore line number pragmas
+                            ln = ln.replace(bad_file_str, good_file_str)
+                            md5.update(ln)
+                    md5.update(b'\n')
+
+        # Select compiler options.
+        compile_opts = ''
+        if os.name == 'nt':
+            compile_opts += '"%s"' % os.path.join(tf.sysconfig.get_lib(), 'python', '_pywrap_tensorflow_internal.lib')
+        elif os.name == 'posix':
+            compile_opts += '"%s"' % os.path.join(tf.sysconfig.get_lib(), 'python', '_pywrap_tensorflow_internal.so')
+            compile_opts += ' --compiler-options \'-fPIC -D_GLIBCXX_USE_CXX11_ABI=0\''
+        else:
+            assert False # not Windows or Linux, w00t?
+        compile_opts += ' --gpu-architecture=%s' % _get_cuda_gpu_arch_string()
+        compile_opts += ' --use_fast_math'
+        nvcc_cmd = _prepare_nvcc_cli(compile_opts)
+
+        # Hash build configuration.
+        md5.update(('nvcc_cmd: ' + nvcc_cmd).encode('utf-8') + b'\n')
+        md5.update(('tf.VERSION: ' + tf.version.VERSION).encode('utf-8') + b'\n')
+        md5.update(('cuda_cache_version_tag: ' + cuda_cache_version_tag).encode('utf-8') + b'\n')
+
+        # Compile if not already compiled.
+        bin_file_ext = '.dll' if os.name == 'nt' else '.so'
+        bin_file = os.path.join(cuda_cache_path, cuda_file_name + '_' + md5.hexdigest() + bin_file_ext)
+        if not os.path.isfile(bin_file):
+            if verbose:
+                print('Compiling... ', end='', flush=True)
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                tmp_file = os.path.join(tmp_dir, cuda_file_name + '_tmp' + bin_file_ext)
+                _run_cmd(nvcc_cmd + ' "%s" --shared -o "%s" --keep --keep-dir "%s"' % (cuda_file, tmp_file, tmp_dir))
+                os.makedirs(cuda_cache_path, exist_ok=True)
+                intermediate_file = os.path.join(cuda_cache_path, cuda_file_name + '_' + uuid.uuid4().hex + '_tmp' + bin_file_ext)
+                shutil.copyfile(tmp_file, intermediate_file)
+                os.rename(intermediate_file, bin_file) # atomic
+
+        # Load.
+        if verbose:
+            print('Loading... ', end='', flush=True)
+        plugin = tf.load_op_library(bin_file)
+
+        # Add to cache.
+        _plugin_cache[cuda_file] = plugin
+        if verbose:
+            print('Done.', flush=True)
+        return plugin
+
+    except:
+        if verbose:
+            print('Failed!', flush=True)
+        raise
+
+#----------------------------------------------------------------------------

stylegan2/layers/cuda/fused_bias_act.cu

@@ -0,0 +1,188 @@
+// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+//
+// This work is made available under the Nvidia Source Code License-NC.
+// To view a copy of this license, visit
+// https://nvlabs.github.io/stylegan2/license.html
+
+#define EIGEN_USE_GPU
+#define __CUDA_INCLUDE_COMPILER_INTERNAL_HEADERS__
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include <stdio.h>
+
+using namespace tensorflow;
+using namespace tensorflow::shape_inference;
+
+#define OP_CHECK_CUDA_ERROR(CTX, CUDA_CALL) do { cudaError_t err = CUDA_CALL; OP_REQUIRES(CTX, err == cudaSuccess, errors::Internal(cudaGetErrorName(err))); } while (false)
+
+//------------------------------------------------------------------------
+// CUDA kernel.
+
+template <class T>
+struct FusedBiasActKernelParams
+{
+    const T*    x;      // [sizeX]
+    const T*    b;      // [sizeB] or NULL
+    const T*    ref;    // [sizeX] or NULL
+    T*          y;      // [sizeX]
+
+    int         grad;
+    int         axis;
+    int         act;
+    float       alpha;
+    float       gain;
+
+    int         sizeX;
+    int         sizeB;
+    int         stepB;
+    int         loopX;
+};
+
+template <class T>
+static __global__ void FusedBiasActKernel(const FusedBiasActKernelParams<T> p)
+{
+    const float expRange        = 80.0f;
+    const float halfExpRange    = 40.0f;
+    const float seluScale       = 1.0507009873554804934193349852946f;
+    const float seluAlpha       = 1.6732632423543772848170429916717f;
+
+    // Loop over elements.
+    int xi = blockIdx.x * p.loopX * blockDim.x + threadIdx.x;
+    for (int loopIdx = 0; loopIdx < p.loopX && xi < p.sizeX; loopIdx++, xi += blockDim.x)
+    {
+        // Load and apply bias.
+        float x = (float)p.x[xi];
+        if (p.b)
+            x += (float)p.b[(xi / p.stepB) % p.sizeB];
+        float ref = (p.ref) ? (float)p.ref[xi] : 0.0f;
+        if (p.gain != 0.0f & p.act != 9)
+            ref /= p.gain;
+
+        // Evaluate activation func.
+        float y;
+        switch (p.act * 10 + p.grad)
+        {
+            // linear
+            default:
+            case 10: y = x; break;
+            case 11: y = x; break;
+            case 12: y = 0.0f; break;
+
+            // relu
+            case 20: y = (x > 0.0f) ? x : 0.0f; break;
+            case 21: y = (ref > 0.0f) ? x : 0.0f; break;
+            case 22: y = 0.0f; break;
+
+            // lrelu
+            case 30: y = (x > 0.0f) ? x : x * p.alpha; break;
+            case 31: y = (ref > 0.0f) ? x : x * p.alpha; break;
+            case 32: y = 0.0f; break;
+
+            // tanh
+            case 40: { float c = expf(x); float d = 1.0f / c; y = (x < -expRange) ? -1.0f : (x > expRange) ? 1.0f : (c - d) / (c + d); } break;
+            case 41: y = x * (1.0f - ref * ref); break;
+            case 42: y = x * (1.0f - ref * ref) * (-2.0f * ref); break;
+
+            // sigmoid
+            case 50: y = (x < -expRange) ? 0.0f : 1.0f / (expf(-x) + 1.0f); break;
+            case 51: y = x * ref * (1.0f - ref); break;
+            case 52: y = x * ref * (1.0f - ref) * (1.0f - 2.0f * ref); break;
+
+            // elu
+            case 60: y = (x >= 0.0f) ? x : expf(x) - 1.0f; break;
+            case 61: y = (ref >= 0.0f) ? x : x * (ref + 1.0f); break;
+            case 62: y = (ref >= 0.0f) ? 0.0f : x * (ref + 1.0f); break;
+
+            // selu
+            case 70: y = (x >= 0.0f) ? seluScale * x : (seluScale * seluAlpha) * (expf(x) - 1.0f); break;
+            case 71: y = (ref >= 0.0f) ? x * seluScale : x * (ref + seluScale * seluAlpha); break;
+            case 72: y = (ref >= 0.0f) ? 0.0f : x * (ref + seluScale * seluAlpha); break;
+
+            // softplus
+            case 80: y = (x > expRange) ? x : logf(expf(x) + 1.0f); break;
+            case 81: y = x * (1.0f - expf(-ref)); break;
+            case 82: { float c = expf(-ref); y = x * c * (1.0f - c); } break;
+
+            // swish
+            case 90: y = (x < -expRange) ? 0.0f : x / (expf(-x) + 1.0f); break;
+            case 91: { float c = expf(ref); float d = c + 1.0f; y = (ref > halfExpRange) ? x : x * c * (ref + d) / (d * d); } break;
+            case 92: { float c = expf(ref); float d = c + 1.0f; y = (ref > halfExpRange) ? 0.0f : x * c * (ref * (2.0f - d) + 2.0f * d) / (d * d * d); } break;
+        }
+
+        // Apply gain and store.
+        p.y[xi] = (T)(y * p.gain);
+    }
+}
+
+//------------------------------------------------------------------------
+// TensorFlow op.
+
+template <class T>
+struct FusedBiasActOp : public OpKernel
+{
+    FusedBiasActKernelParams<T> m_attribs;
+
+    FusedBiasActOp(OpKernelConstruction* ctx) : OpKernel(ctx)
+    {
+        memset(&m_attribs, 0, sizeof(m_attribs));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("grad", &m_attribs.grad));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("axis", &m_attribs.axis));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("act", &m_attribs.act));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("alpha", &m_attribs.alpha));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("gain", &m_attribs.gain));
+        OP_REQUIRES(ctx, m_attribs.grad >= 0, errors::InvalidArgument("grad must be non-negative"));
+        OP_REQUIRES(ctx, m_attribs.axis >= 0, errors::InvalidArgument("axis must be non-negative"));
+        OP_REQUIRES(ctx, m_attribs.act >= 0, errors::InvalidArgument("act must be non-negative"));
+    }
+
+    void Compute(OpKernelContext* ctx)
+    {
+        FusedBiasActKernelParams<T> p = m_attribs;
+        cudaStream_t stream = ctx->eigen_device<Eigen::GpuDevice>().stream();
+
+        const Tensor& x     = ctx->input(0); // [...]
+        const Tensor& b     = ctx->input(1); // [sizeB] or [0]
+        const Tensor& ref   = ctx->input(2); // x.shape or [0]
+        p.x = x.flat<T>().data();
+        p.b = (b.NumElements()) ? b.flat<T>().data() : NULL;
+        p.ref = (ref.NumElements()) ? ref.flat<T>().data() : NULL;
+        OP_REQUIRES(ctx, b.NumElements() == 0 || m_attribs.axis < x.dims(), errors::InvalidArgument("axis out of bounds"));
+        OP_REQUIRES(ctx, b.dims() == 1, errors::InvalidArgument("b must have rank 1"));
+        OP_REQUIRES(ctx, b.NumElements() == 0 || b.NumElements() == x.dim_size(m_attribs.axis), errors::InvalidArgument("b has wrong number of elements"));
+        OP_REQUIRES(ctx, ref.NumElements() == ((p.grad == 0) ? 0 : x.NumElements()), errors::InvalidArgument("ref has wrong number of elements"));
+        OP_REQUIRES(ctx, x.NumElements() <= kint32max, errors::InvalidArgument("x is too large"));
+
+        p.sizeX = (int)x.NumElements();
+        p.sizeB = (int)b.NumElements();
+        p.stepB = 1;
+        for (int i = m_attribs.axis + 1; i < x.dims(); i++)
+            p.stepB *= (int)x.dim_size(i);
+
+        Tensor* y = NULL; // x.shape
+        OP_REQUIRES_OK(ctx, ctx->allocate_output(0, x.shape(), &y));
+        p.y = y->flat<T>().data();
+
+        p.loopX = 4;
+        int blockSize = 4 * 32;
+        int gridSize = (p.sizeX - 1) / (p.loopX * blockSize) + 1;
+        void* args[] = {&p};
+        OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel((void*)FusedBiasActKernel<T>, gridSize, blockSize, args, 0, stream));
+    }
+};
+
+REGISTER_OP("FusedBiasAct")
+    .Input      ("x: T")
+    .Input      ("b: T")
+    .Input      ("ref: T")
+    .Output     ("y: T")
+    .Attr       ("T: {float, half}")
+    .Attr       ("grad: int = 0")
+    .Attr       ("axis: int = 1")
+    .Attr       ("act: int = 0")
+    .Attr       ("alpha: float = 0.0")
+    .Attr       ("gain: float = 1.0");
+REGISTER_KERNEL_BUILDER(Name("FusedBiasAct").Device(DEVICE_GPU).TypeConstraint<float>("T"), FusedBiasActOp<float>);
+REGISTER_KERNEL_BUILDER(Name("FusedBiasAct").Device(DEVICE_GPU).TypeConstraint<Eigen::half>("T"), FusedBiasActOp<Eigen::half>);
+
+//------------------------------------------------------------------------

stylegan2/layers/cuda/fused_bias_act.py

@@ -0,0 +1,196 @@
+# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://nvlabs.github.io/stylegan2/license.html
+
+"""Custom TensorFlow ops for efficient bias and activation."""
+
+import os
+import numpy as np
+import tensorflow as tf
+from stylegan2.layers.cuda import custom_ops
+from stylegan2.utils import EasyDict
+
+def _get_plugin():
+    return custom_ops.get_plugin(os.path.splitext(__file__)[0] + '.cu')
+
+#----------------------------------------------------------------------------
+
+activation_funcs = {
+    'linear':   EasyDict(func=lambda x, **_:        x,                          def_alpha=None, def_gain=1.0,           cuda_idx=1, ref='y', zero_2nd_grad=True),
+    'relu':     EasyDict(func=lambda x, **_:        tf.nn.relu(x),              def_alpha=None, def_gain=np.sqrt(2),    cuda_idx=2, ref='y', zero_2nd_grad=True),
+    'lrelu':    EasyDict(func=lambda x, alpha, **_: tf.nn.leaky_relu(x, alpha), def_alpha=0.2,  def_gain=np.sqrt(2),    cuda_idx=3, ref='y', zero_2nd_grad=True),
+    'tanh':     EasyDict(func=lambda x, **_:        tf.nn.tanh(x),              def_alpha=None, def_gain=1.0,           cuda_idx=4, ref='y', zero_2nd_grad=False),
+    'sigmoid':  EasyDict(func=lambda x, **_:        tf.nn.sigmoid(x),           def_alpha=None, def_gain=1.0,           cuda_idx=5, ref='y', zero_2nd_grad=False),
+    'elu':      EasyDict(func=lambda x, **_:        tf.nn.elu(x),               def_alpha=None, def_gain=1.0,           cuda_idx=6, ref='y', zero_2nd_grad=False),
+    'selu':     EasyDict(func=lambda x, **_:        tf.nn.selu(x),              def_alpha=None, def_gain=1.0,           cuda_idx=7, ref='y', zero_2nd_grad=False),
+    'softplus': EasyDict(func=lambda x, **_:        tf.nn.softplus(x),          def_alpha=None, def_gain=1.0,           cuda_idx=8, ref='y', zero_2nd_grad=False),
+    'swish':    EasyDict(func=lambda x, **_:        tf.nn.sigmoid(x) * x,       def_alpha=None, def_gain=np.sqrt(2),    cuda_idx=9, ref='x', zero_2nd_grad=False),
+}
+
+#----------------------------------------------------------------------------
+
+def fused_bias_act(x, b=None, axis=1, act='linear', alpha=None, gain=None, impl='cuda'):
+    r"""Fused bias and activation function.
+
+    Adds bias `b` to activation tensor `x`, evaluates activation function `act`,
+    and scales the result by `gain`. Each of the steps is optional. In most cases,
+    the fused op is considerably more efficient than performing the same calculation
+    using standard TensorFlow ops. It supports first and second order gradients,
+    but not third order gradients.
+
+    Args:
+        x:      Input activation tensor. Can have any shape, but if `b` is defined, the
+                dimension corresponding to `axis`, as well as the rank, must be known.
+        b:      Bias vector, or `None` to disable. Must be a 1D tensor of the same type
+                as `x`. The shape must be known, and it must match the dimension of `x`
+                corresponding to `axis`.
+        axis:   The dimension in `x` corresponding to the elements of `b`.
+                The value of `axis` is ignored if `b` is not specified.
+        act:    Name of the activation function to evaluate, or `"linear"` to disable.
+                Can be e.g. `"relu"`, `"lrelu"`, `"tanh"`, `"sigmoid"`, `"swish"`, etc.
+                See `activation_funcs` for a full list. `None` is not allowed.
+        alpha:  Shape parameter for the activation function, or `None` to use the default.
+        gain:   Scaling factor for the output tensor, or `None` to use default.
+                See `activation_funcs` for the default scaling of each activation function.
+                If unsure, consider specifying `1.0`.
+        impl:   Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
+
+    Returns:
+        Tensor of the same shape and datatype as `x`.
+    """
+
+    impl_dict = {
+        'ref':  _fused_bias_act_ref,
+        'cuda': _fused_bias_act_cuda,
+    }
+    return impl_dict[impl](x=x, b=b, axis=axis, act=act, alpha=alpha, gain=gain)
+
+#----------------------------------------------------------------------------
+
+def _fused_bias_act_ref(x, b, axis, act, alpha, gain):
+    """Slow reference implementation of `fused_bias_act()` using standard TensorFlow ops."""
+
+    # Validate arguments.
+    x = tf.convert_to_tensor(x)
+    b = tf.convert_to_tensor(b) if b is not None else tf.constant([], dtype=x.dtype)
+    act_spec = activation_funcs[act]
+    assert b.shape.rank == 1 and (b.shape[0] == 0 or b.shape[0] == x.shape[axis])
+    assert b.shape[0] == 0 or 0 <= axis < x.shape.rank
+    if alpha is None:
+        alpha = act_spec.def_alpha
+    if gain is None:
+        gain = act_spec.def_gain
+
+    # Add bias.
+    if b.shape[0] != 0:
+        x += tf.reshape(b, [-1 if i == axis else 1 for i in range(x.shape.rank)])
+
+    # Evaluate activation function.
+    x = act_spec.func(x, alpha=alpha)
+
+    # Scale by gain.
+    if gain != 1:
+        x *= gain
+    return x
+
+#----------------------------------------------------------------------------
+
+def _fused_bias_act_cuda(x, b, axis, act, alpha, gain):
+    """Fast CUDA implementation of `fused_bias_act()` using custom ops."""
+
+    # Validate arguments.
+    x = tf.convert_to_tensor(x)
+    empty_tensor = tf.constant([], dtype=x.dtype)
+    b = tf.convert_to_tensor(b) if b is not None else empty_tensor
+    act_spec = activation_funcs[act]
+    assert b.shape.rank == 1 and (b.shape[0] == 0 or b.shape[0] == x.shape[axis])
+    assert b.shape[0] == 0 or 0 <= axis < x.shape.rank
+    if alpha is None:
+        alpha = act_spec.def_alpha
+    if gain is None:
+        gain = act_spec.def_gain
+
+    # Special cases.
+    if act == 'linear' and b is None and gain == 1.0:
+        return x
+    if act_spec.cuda_idx is None:
+        return _fused_bias_act_ref(x=x, b=b, axis=axis, act=act, alpha=alpha, gain=gain)
+
+    # CUDA kernel.
+    cuda_kernel = _get_plugin().fused_bias_act
+    cuda_kwargs = dict(axis=axis, act=act_spec.cuda_idx, alpha=alpha, gain=gain)
+
+    # Forward pass: y = func(x, b).
+    def func_y(x, b):
+        y = cuda_kernel(x=x, b=b, ref=empty_tensor, grad=0, **cuda_kwargs)
+        y.set_shape(x.shape)
+        return y
+
+    # Backward pass: dx, db = grad(dy, x, y)
+    def grad_dx(dy, x, y):
+        ref = {'x': x, 'y': y}[act_spec.ref]
+        dx = cuda_kernel(x=dy, b=empty_tensor, ref=ref, grad=1, **cuda_kwargs)
+        dx.set_shape(x.shape)
+        return dx
+    def grad_db(dx):
+        if b.shape[0] == 0:
+            return empty_tensor
+        db = dx
+        if axis < x.shape.rank - 1:
+            db = tf.reduce_sum(db, list(range(axis + 1, x.shape.rank)))
+        if axis > 0:
+            db = tf.reduce_sum(db, list(range(axis)))
+        db.set_shape(b.shape)
+        return db
+
+    # Second order gradients: d_dy, d_x = grad2(d_dx, d_db, x, y)
+    def grad2_d_dy(d_dx, d_db, x, y):
+        ref = {'x': x, 'y': y}[act_spec.ref]
+        d_dy = cuda_kernel(x=d_dx, b=d_db, ref=ref, grad=1, **cuda_kwargs)
+        d_dy.set_shape(x.shape)
+        return d_dy
+    def grad2_d_x(d_dx, d_db, x, y):
+        ref = {'x': x, 'y': y}[act_spec.ref]
+        d_x = cuda_kernel(x=d_dx, b=d_db, ref=ref, grad=2, **cuda_kwargs)
+        d_x.set_shape(x.shape)
+        return d_x
+
+    # Fast version for piecewise-linear activation funcs.
+    @tf.custom_gradient
+    def func_zero_2nd_grad(x, b):
+        y = func_y(x, b)
+        @tf.custom_gradient
+        def grad(dy):
+            dx = grad_dx(dy, x, y)
+            db = grad_db(dx)
+            def grad2(d_dx, d_db):
+                d_dy = grad2_d_dy(d_dx, d_db, x, y)
+                return d_dy
+            return (dx, db), grad2
+        return y, grad
+
+    # Slow version for general activation funcs.
+    @tf.custom_gradient
+    def func_nonzero_2nd_grad(x, b):
+        y = func_y(x, b)
+        def grad_wrap(dy):
+            @tf.custom_gradient
+            def grad_impl(dy, x):
+                dx = grad_dx(dy, x, y)
+                db = grad_db(dx)
+                def grad2(d_dx, d_db):
+                    d_dy = grad2_d_dy(d_dx, d_db, x, y)
+                    d_x = grad2_d_x(d_dx, d_db, x, y)
+                    return d_dy, d_x
+                return (dx, db), grad2
+            return grad_impl(dy, x)
+        return y, grad_wrap
+
+    # Which version to use?
+    if act_spec.zero_2nd_grad:
+        return func_zero_2nd_grad(x, b)
+    return func_nonzero_2nd_grad(x, b)
+
+#----------------------------------------------------------------------------

stylegan2/layers/cuda/upfirdn_2d.cu

@@ -0,0 +1,326 @@
+// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+//
+// This work is made available under the Nvidia Source Code License-NC.
+// To view a copy of this license, visit
+// https://nvlabs.github.io/stylegan2/license.html
+
+#define EIGEN_USE_GPU
+#define __CUDA_INCLUDE_COMPILER_INTERNAL_HEADERS__
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include <stdio.h>
+
+using namespace tensorflow;
+using namespace tensorflow::shape_inference;
+
+//------------------------------------------------------------------------
+// Helpers.
+
+#define OP_CHECK_CUDA_ERROR(CTX, CUDA_CALL) do { cudaError_t err = CUDA_CALL; OP_REQUIRES(CTX, err == cudaSuccess, errors::Internal(cudaGetErrorName(err))); } while (false)
+
+static __host__ __device__ __forceinline__ int floorDiv(int a, int b)
+{
+    int c = a / b;
+    if (c * b > a)
+        c--;
+    return c;
+}
+
+//------------------------------------------------------------------------
+// CUDA kernel params.
+
+template <class T>
+struct UpFirDn2DKernelParams
+{
+    const T*    x;          // [majorDim, inH, inW, minorDim]
+    const T*    k;          // [kernelH, kernelW]
+    T*          y;          // [majorDim, outH, outW, minorDim]
+
+    int         upx;
+    int         upy;
+    int         downx;
+    int         downy;
+    int         padx0;
+    int         padx1;
+    int         pady0;
+    int         pady1;
+
+    int         majorDim;
+    int         inH;
+    int         inW;
+    int         minorDim;
+    int         kernelH;
+    int         kernelW;
+    int         outH;
+    int         outW;
+    int         loopMajor;
+    int         loopX;
+};
+
+//------------------------------------------------------------------------
+// General CUDA implementation for large filter kernels.
+
+template <class T>
+static __global__ void UpFirDn2DKernel_large(const UpFirDn2DKernelParams<T> p)
+{
+    // Calculate thread index.
+    int minorIdx = blockIdx.x * blockDim.x + threadIdx.x;
+    int outY = minorIdx / p.minorDim;
+    minorIdx -= outY * p.minorDim;
+    int outXBase = blockIdx.y * p.loopX * blockDim.y + threadIdx.y;
+    int majorIdxBase = blockIdx.z * p.loopMajor;
+    if (outXBase >= p.outW || outY >= p.outH || majorIdxBase >= p.majorDim)
+        return;
+
+    // Setup Y receptive field.
+    int midY = outY * p.downy + p.upy - 1 - p.pady0;
+    int inY = min(max(floorDiv(midY, p.upy), 0), p.inH);
+    int h = min(max(floorDiv(midY + p.kernelH, p.upy), 0), p.inH) - inY;
+    int kernelY = midY + p.kernelH - (inY + 1) * p.upy;
+
+    // Loop over majorDim and outX.
+    for (int loopMajor = 0, majorIdx = majorIdxBase; loopMajor < p.loopMajor && majorIdx < p.majorDim; loopMajor++, majorIdx++)
+    for (int loopX = 0, outX = outXBase; loopX < p.loopX && outX < p.outW; loopX++, outX += blockDim.y)
+    {
+        // Setup X receptive field.
+        int midX = outX * p.downx + p.upx - 1 - p.padx0;
+        int inX = min(max(floorDiv(midX, p.upx), 0), p.inW);
+        int w = min(max(floorDiv(midX + p.kernelW, p.upx), 0), p.inW) - inX;
+        int kernelX = midX + p.kernelW - (inX + 1) * p.upx;
+
+        // Initialize pointers.
+        const T* xp = &p.x[((majorIdx * p.inH + inY) * p.inW + inX) * p.minorDim + minorIdx];
+        const T* kp = &p.k[kernelY * p.kernelW + kernelX];
+        int xpx = p.minorDim;
+        int kpx = -p.upx;
+        int xpy = p.inW * p.minorDim;
+        int kpy = -p.upy * p.kernelW;
+
+        // Inner loop.
+        float v = 0.0f;
+        for (int y = 0; y < h; y++)
+        {
+            for (int x = 0; x < w; x++)
+            {
+                v += (float)(*xp) * (float)(*kp);
+                xp += xpx;
+                kp += kpx;
+            }
+            xp += xpy - w * xpx;
+            kp += kpy - w * kpx;
+        }
+
+        // Store result.
+        p.y[((majorIdx * p.outH + outY) * p.outW + outX) * p.minorDim + minorIdx] = (T)v;
+    }
+}
+
+//------------------------------------------------------------------------
+// Specialized CUDA implementation for small filter kernels.
+
+template <class T, int upx, int upy, int downx, int downy, int kernelW, int kernelH, int tileOutW, int tileOutH>
+static __global__ void UpFirDn2DKernel_small(const UpFirDn2DKernelParams<T> p)
+{
+    //assert(kernelW % upx == 0);
+    //assert(kernelH % upy == 0);
+    const int tileInW = ((tileOutW - 1) * downx + kernelW - 1) / upx + 1;
+    const int tileInH = ((tileOutH - 1) * downy + kernelH - 1) / upy + 1;
+    __shared__ volatile float sk[kernelH][kernelW];
+    __shared__ volatile float sx[tileInH][tileInW];
+
+    // Calculate tile index.
+    int minorIdx = blockIdx.x;
+    int tileOutY = minorIdx / p.minorDim;
+    minorIdx -= tileOutY * p.minorDim;
+    tileOutY *= tileOutH;
+    int tileOutXBase = blockIdx.y * p.loopX * tileOutW;
+    int majorIdxBase = blockIdx.z * p.loopMajor;
+    if (tileOutXBase >= p.outW | tileOutY >= p.outH | majorIdxBase >= p.majorDim)
+        return;
+
+    // Load filter kernel (flipped).
+    for (int tapIdx = threadIdx.x; tapIdx < kernelH * kernelW; tapIdx += blockDim.x)
+    {
+        int ky = tapIdx / kernelW;
+        int kx = tapIdx - ky * kernelW;
+        float v = 0.0f;
+        if (kx < p.kernelW & ky < p.kernelH)
+            v = (float)p.k[(p.kernelH - 1 - ky) * p.kernelW + (p.kernelW - 1 - kx)];
+        sk[ky][kx] = v;
+    }
+
+    // Loop over majorDim and outX.
+    for (int loopMajor = 0, majorIdx = majorIdxBase; loopMajor < p.loopMajor & majorIdx < p.majorDim; loopMajor++, majorIdx++)
+    for (int loopX = 0, tileOutX = tileOutXBase; loopX < p.loopX & tileOutX < p.outW; loopX++, tileOutX += tileOutW)
+    {
+        // Load input pixels.
+        int tileMidX = tileOutX * downx + upx - 1 - p.padx0;
+        int tileMidY = tileOutY * downy + upy - 1 - p.pady0;
+        int tileInX = floorDiv(tileMidX, upx);
+        int tileInY = floorDiv(tileMidY, upy);
+        __syncthreads();
+        for (int inIdx = threadIdx.x; inIdx < tileInH * tileInW; inIdx += blockDim.x)
+        {
+            int relInY = inIdx / tileInW;
+            int relInX = inIdx - relInY * tileInW;
+            int inX = relInX + tileInX;
+            int inY = relInY + tileInY;
+            float v = 0.0f;
+            if (inX >= 0 & inY >= 0 & inX < p.inW & inY < p.inH)
+                v = (float)p.x[((majorIdx * p.inH + inY) * p.inW + inX) * p.minorDim + minorIdx];
+            sx[relInY][relInX] = v;
+        }
+
+        // Loop over output pixels.
+        __syncthreads();
+        for (int outIdx = threadIdx.x; outIdx < tileOutH * tileOutW; outIdx += blockDim.x)
+        {
+            int relOutY = outIdx / tileOutW;
+            int relOutX = outIdx - relOutY * tileOutW;
+            int outX = relOutX + tileOutX;
+            int outY = relOutY + tileOutY;
+
+            // Setup receptive field.
+            int midX = tileMidX + relOutX * downx;
+            int midY = tileMidY + relOutY * downy;
+            int inX = floorDiv(midX, upx);
+            int inY = floorDiv(midY, upy);
+            int relInX = inX - tileInX;
+            int relInY = inY - tileInY;
+            int kernelX = (inX + 1) * upx - midX - 1; // flipped
+            int kernelY = (inY + 1) * upy - midY - 1; // flipped
+
+            // Inner loop.
+            float v = 0.0f;
+            #pragma unroll
+            for (int y = 0; y < kernelH / upy; y++)
+                #pragma unroll
+                for (int x = 0; x < kernelW / upx; x++)
+                    v += sx[relInY + y][relInX + x] * sk[kernelY + y * upy][kernelX + x * upx];
+
+            // Store result.
+            if (outX < p.outW & outY < p.outH)
+                p.y[((majorIdx * p.outH + outY) * p.outW + outX) * p.minorDim + minorIdx] = (T)v;
+        }
+    }
+}
+
+//------------------------------------------------------------------------
+// TensorFlow op.
+
+template <class T>
+struct UpFirDn2DOp : public OpKernel
+{
+    UpFirDn2DKernelParams<T> m_attribs;
+
+    UpFirDn2DOp(OpKernelConstruction* ctx) : OpKernel(ctx)
+    {
+        memset(&m_attribs, 0, sizeof(m_attribs));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("upx", &m_attribs.upx));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("upy", &m_attribs.upy));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("downx", &m_attribs.downx));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("downy", &m_attribs.downy));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("padx0", &m_attribs.padx0));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("padx1", &m_attribs.padx1));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("pady0", &m_attribs.pady0));
+        OP_REQUIRES_OK(ctx, ctx->GetAttr("pady1", &m_attribs.pady1));
+        OP_REQUIRES(ctx, m_attribs.upx >= 1 && m_attribs.upy >= 1, errors::InvalidArgument("upx and upy must be at least 1x1"));
+        OP_REQUIRES(ctx, m_attribs.downx >= 1 && m_attribs.downy >= 1, errors::InvalidArgument("downx and downy must be at least 1x1"));
+    }
+
+    void Compute(OpKernelContext* ctx)
+    {
+        UpFirDn2DKernelParams<T> p = m_attribs;
+        cudaStream_t stream = ctx->eigen_device<Eigen::GpuDevice>().stream();
+
+        const Tensor& x = ctx->input(0); // [majorDim, inH, inW, minorDim]
+        const Tensor& k = ctx->input(1); // [kernelH, kernelW]
+        p.x = x.flat<T>().data();
+        p.k = k.flat<T>().data();
+        OP_REQUIRES(ctx, x.dims() == 4, errors::InvalidArgument("input must have rank 4"));
+        OP_REQUIRES(ctx, k.dims() == 2, errors::InvalidArgument("kernel must have rank 2"));
+        OP_REQUIRES(ctx, x.NumElements() <= kint32max, errors::InvalidArgument("input too large"));
+        OP_REQUIRES(ctx, k.NumElements() <= kint32max, errors::InvalidArgument("kernel too large"));
+
+        p.majorDim  = (int)x.dim_size(0);
+        p.inH       = (int)x.dim_size(1);
+        p.inW       = (int)x.dim_size(2);
+        p.minorDim  = (int)x.dim_size(3);
+        p.kernelH   = (int)k.dim_size(0);
+        p.kernelW   = (int)k.dim_size(1);
+        OP_REQUIRES(ctx, p.kernelW >= 1 && p.kernelH >= 1, errors::InvalidArgument("kernel must be at least 1x1"));
+
+        p.outW = (p.inW * p.upx + p.padx0 + p.padx1 - p.kernelW + p.downx) / p.downx;
+        p.outH = (p.inH * p.upy + p.pady0 + p.pady1 - p.kernelH + p.downy) / p.downy;
+        OP_REQUIRES(ctx, p.outW >= 1 && p.outH >= 1, errors::InvalidArgument("output must be at least 1x1"));
+
+        Tensor* y = NULL; // [majorDim, outH, outW, minorDim]
+        TensorShape ys;
+        ys.AddDim(p.majorDim);
+        ys.AddDim(p.outH);
+        ys.AddDim(p.outW);
+        ys.AddDim(p.minorDim);
+        OP_REQUIRES_OK(ctx, ctx->allocate_output(0, ys, &y));
+        p.y = y->flat<T>().data();
+        OP_REQUIRES(ctx, y->NumElements() <= kint32max, errors::InvalidArgument("output too large"));
+
+        // Choose CUDA kernel to use.
+        void* cudaKernel = (void*)UpFirDn2DKernel_large<T>;
+        int tileOutW = -1;
+        int tileOutH = -1;
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 7 && p.kernelH <= 7) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 7,7, 64,16>; tileOutW = 64; tileOutH = 16; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 6 && p.kernelH <= 6) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 6,6, 64,16>; tileOutW = 64; tileOutH = 16; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 5 && p.kernelH <= 5) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 5,5, 64,16>; tileOutW = 64; tileOutH = 16; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 4 && p.kernelH <= 4) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 4,4, 64,16>; tileOutW = 64; tileOutH = 16; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 1 && p.downy == 1 && p.kernelW <= 3 && p.kernelH <= 3) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 1,1, 3,3, 64,16>; tileOutW = 64; tileOutH = 16; }
+        if (p.upx == 2 && p.upy == 2 && p.downx == 1 && p.downy == 1 && p.kernelW <= 8 && p.kernelH <= 8) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 2,2, 1,1, 8,8, 64,16>; tileOutW = 64; tileOutH = 16; }
+        if (p.upx == 2 && p.upy == 2 && p.downx == 1 && p.downy == 1 && p.kernelW <= 6 && p.kernelH <= 6) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 2,2, 1,1, 6,6, 64,16>; tileOutW = 64; tileOutH = 16; }
+        if (p.upx == 2 && p.upy == 2 && p.downx == 1 && p.downy == 1 && p.kernelW <= 4 && p.kernelH <= 4) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 2,2, 1,1, 4,4, 64,16>; tileOutW = 64; tileOutH = 16; }
+        if (p.upx == 2 && p.upy == 2 && p.downx == 1 && p.downy == 1 && p.kernelW <= 2 && p.kernelH <= 2) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 2,2, 1,1, 2,2, 64,16>; tileOutW = 64; tileOutH = 16; }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 2 && p.downy == 2 && p.kernelW <= 8 && p.kernelH <= 8) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 2,2, 8,8, 32,8>;  tileOutW = 32; tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 2 && p.downy == 2 && p.kernelW <= 6 && p.kernelH <= 6) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 2,2, 6,6, 32,8>;  tileOutW = 32; tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 2 && p.downy == 2 && p.kernelW <= 4 && p.kernelH <= 4) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 2,2, 4,4, 32,8>;  tileOutW = 32; tileOutH = 8;  }
+        if (p.upx == 1 && p.upy == 1 && p.downx == 2 && p.downy == 2 && p.kernelW <= 2 && p.kernelH <= 2) { cudaKernel = (void*)UpFirDn2DKernel_small<T, 1,1, 2,2, 2,2, 32,8>;  tileOutW = 32; tileOutH = 8;  }
+
+        // Choose launch params.
+        dim3 blockSize;
+        dim3 gridSize;
+        if (tileOutW > 0 && tileOutH > 0) // small
+        {
+            p.loopMajor = (p.majorDim - 1) / 16384 + 1;
+            p.loopX = 1;
+            blockSize = dim3(32 * 8, 1, 1);
+            gridSize = dim3(((p.outH - 1) / tileOutH + 1) * p.minorDim, (p.outW - 1) / (p.loopX * tileOutW) + 1, (p.majorDim - 1) / p.loopMajor + 1);
+        }
+        else // large
+        {
+            p.loopMajor = (p.majorDim - 1) / 16384 + 1;
+            p.loopX = 4;
+            blockSize = dim3(4, 32, 1);
+            gridSize = dim3((p.outH * p.minorDim - 1) / blockSize.x + 1, (p.outW - 1) / (p.loopX * blockSize.y) + 1, (p.majorDim - 1) / p.loopMajor + 1);
+        }
+
+        // Launch CUDA kernel.
+        void* args[] = {&p};
+        OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel(cudaKernel, gridSize, blockSize, args, 0, stream));
+    }
+};
+
+REGISTER_OP("UpFirDn2D")
+    .Input      ("x: T")
+    .Input      ("k: T")
+    .Output     ("y: T")
+    .Attr       ("T: {float, half}")
+    .Attr       ("upx: int = 1")
+    .Attr       ("upy: int = 1")
+    .Attr       ("downx: int = 1")
+    .Attr       ("downy: int = 1")
+    .Attr       ("padx0: int = 0")
+    .Attr       ("padx1: int = 0")
+    .Attr       ("pady0: int = 0")
+    .Attr       ("pady1: int = 0");
+REGISTER_KERNEL_BUILDER(Name("UpFirDn2D").Device(DEVICE_GPU).TypeConstraint<float>("T"), UpFirDn2DOp<float>);
+REGISTER_KERNEL_BUILDER(Name("UpFirDn2D").Device(DEVICE_GPU).TypeConstraint<Eigen::half>("T"), UpFirDn2DOp<Eigen::half>);
+
+//------------------------------------------------------------------------

stylegan2/layers/cuda/upfirdn_2d.py

@@ -0,0 +1,364 @@
+# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://nvlabs.github.io/stylegan2/license.html
+
+"""Custom TensorFlow ops for efficient resampling of 2D images."""
+
+import os
+import numpy as np
+import tensorflow as tf
+from stylegan2.layers.cuda import custom_ops
+
+def _get_plugin():
+    return custom_ops.get_plugin(os.path.splitext(__file__)[0] + '.cu')
+
+#----------------------------------------------------------------------------
+
+def upfirdn_2d(x, k, upx=1, upy=1, downx=1, downy=1, padx0=0, padx1=0, pady0=0, pady1=0, impl='cuda'):
+    r"""Pad, upsample, FIR filter, and downsample a batch of 2D images.
+
+    Accepts a batch of 2D images of the shape `[majorDim, inH, inW, minorDim]`
+    and performs the following operations for each image, batched across
+    `majorDim` and `minorDim`:
+
+    1. Pad the image with zeros by the specified number of pixels on each side
+       (`padx0`, `padx1`, `pady0`, `pady1`). Specifying a negative value
+       corresponds to cropping the image.
+
+    2. Upsample the image by inserting the zeros after each pixel (`upx`, `upy`).
+
+    3. Convolve the image with the specified 2D FIR filter (`k`), shrinking the
+       image so that the footprint of all output pixels lies within the input image.
+
+    4. Downsample the image by throwing away pixels (`downx`, `downy`).
+
+    This sequence of operations bears close resemblance to scipy.signal.upfirdn().
+    The fused op is considerably more efficient than performing the same calculation
+    using standard TensorFlow ops. It supports gradients of arbitrary order.
+
+    Args:
+        x:      Input tensor of the shape `[majorDim, inH, inW, minorDim]`.
+        k:      2D FIR filter of the shape `[firH, firW]`.
+        upx:    Integer upsampling factor along the X-axis (default: 1).
+        upy:    Integer upsampling factor along the Y-axis (default: 1).
+        downx:  Integer downsampling factor along the X-axis (default: 1).
+        downy:  Integer downsampling factor along the Y-axis (default: 1).
+        padx0:  Number of pixels to pad on the left side (default: 0).
+        padx1:  Number of pixels to pad on the right side (default: 0).
+        pady0:  Number of pixels to pad on the top side (default: 0).
+        pady1:  Number of pixels to pad on the bottom side (default: 0).
+        impl:   Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
+
+    Returns:
+        Tensor of the shape `[majorDim, outH, outW, minorDim]`, and same datatype as `x`.
+    """
+
+    impl_dict = {
+        'ref':  _upfirdn_2d_ref,
+        'cuda': _upfirdn_2d_cuda,
+    }
+    return impl_dict[impl](x=x, k=k, upx=upx, upy=upy, downx=downx, downy=downy, padx0=padx0, padx1=padx1, pady0=pady0, pady1=pady1)
+
+#----------------------------------------------------------------------------
+
+def _upfirdn_2d_ref(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1):
+    """Slow reference implementation of `upfirdn_2d()` using standard TensorFlow ops."""
+
+    x = tf.convert_to_tensor(x)
+    k = np.asarray(k, dtype=np.float32)
+    assert x.shape.rank == 4
+    inH = x.shape[1]
+    inW = x.shape[2]
+    minorDim = _shape(x, 3)
+    kernelH, kernelW = k.shape
+    assert inW >= 1 and inH >= 1
+    assert kernelW >= 1 and kernelH >= 1
+    assert isinstance(upx, int) and isinstance(upy, int)
+    assert isinstance(downx, int) and isinstance(downy, int)
+    assert isinstance(padx0, int) and isinstance(padx1, int)
+    assert isinstance(pady0, int) and isinstance(pady1, int)
+
+    # Upsample (insert zeros).
+    x = tf.reshape(x, [-1, inH, 1, inW, 1, minorDim])
+    x = tf.pad(x, [[0, 0], [0, 0], [0, upy - 1], [0, 0], [0, upx - 1], [0, 0]])
+    x = tf.reshape(x, [-1, inH * upy, inW * upx, minorDim])
+
+    # Pad (crop if negative).
+    x = tf.pad(x, [[0, 0], [max(pady0, 0), max(pady1, 0)], [max(padx0, 0), max(padx1, 0)], [0, 0]])
+    x = x[:, max(-pady0, 0) : x.shape[1] - max(-pady1, 0), max(-padx0, 0) : x.shape[2] - max(-padx1, 0), :]
+
+    # Convolve with filter.
+    x = tf.transpose(x, [0, 3, 1, 2])
+    x = tf.reshape(x, [-1, 1, inH * upy + pady0 + pady1, inW * upx + padx0 + padx1])
+    w = tf.constant(k[::-1, ::-1, np.newaxis, np.newaxis], dtype=x.dtype)
+    x = tf.nn.conv2d(x, w, strides=[1,1,1,1], padding='VALID', data_format='NCHW')
+    x = tf.reshape(x, [-1, minorDim, inH * upy + pady0 + pady1 - kernelH + 1, inW * upx + padx0 + padx1 - kernelW + 1])
+    x = tf.transpose(x, [0, 2, 3, 1])
+
+    # Downsample (throw away pixels).
+    return x[:, ::downy, ::downx, :]
+
+#----------------------------------------------------------------------------
+
+def _upfirdn_2d_cuda(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1):
+    """Fast CUDA implementation of `upfirdn_2d()` using custom ops."""
+
+    x = tf.convert_to_tensor(x)
+    k = np.asarray(k, dtype=np.float32)
+    majorDim, inH, inW, minorDim = x.shape.as_list()
+    kernelH, kernelW = k.shape
+    assert inW >= 1 and inH >= 1
+    assert kernelW >= 1 and kernelH >= 1
+    assert isinstance(upx, int) and isinstance(upy, int)
+    assert isinstance(downx, int) and isinstance(downy, int)
+    assert isinstance(padx0, int) and isinstance(padx1, int)
+    assert isinstance(pady0, int) and isinstance(pady1, int)
+
+    outW = (inW * upx + padx0 + padx1 - kernelW) // downx + 1
+    outH = (inH * upy + pady0 + pady1 - kernelH) // downy + 1
+    assert outW >= 1 and outH >= 1
+
+    kc = tf.constant(k, dtype=x.dtype)
+    gkc = tf.constant(k[::-1, ::-1], dtype=x.dtype)
+    gpadx0 = kernelW - padx0 - 1
+    gpady0 = kernelH - pady0 - 1
+    gpadx1 = inW * upx - outW * downx + padx0 - upx + 1
+    gpady1 = inH * upy - outH * downy + pady0 - upy + 1
+
+    @tf.custom_gradient
+    def func(x):
+        y = _get_plugin().up_fir_dn2d(x=x, k=kc, upx=upx, upy=upy, downx=downx, downy=downy, padx0=padx0, padx1=padx1, pady0=pady0, pady1=pady1)
+        y.set_shape([majorDim, outH, outW, minorDim])
+        @tf.custom_gradient
+        def grad(dy):
+            dx = _get_plugin().up_fir_dn2d(x=dy, k=gkc, upx=downx, upy=downy, downx=upx, downy=upy, padx0=gpadx0, padx1=gpadx1, pady0=gpady0, pady1=gpady1)
+            dx.set_shape([majorDim, inH, inW, minorDim])
+            return dx, func
+        return y, grad
+    return func(x)
+
+#----------------------------------------------------------------------------
+
+def filter_2d(x, k, gain=1, data_format='NCHW', impl='cuda'):
+    r"""Filter a batch of 2D images with the given FIR filter.
+
+    Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]`
+    and filters each image with the given filter. The filter is normalized so that
+    if the input pixels are constant, they will be scaled by the specified `gain`.
+    Pixels outside the image are assumed to be zero.
+
+    Args:
+        x:            Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.
+        k:            FIR filter of the shape `[firH, firW]` or `[firN]` (separable).
+        gain:         Scaling factor for signal magnitude (default: 1.0).
+        data_format:  `'NCHW'` or `'NHWC'` (default: `'NCHW'`).
+        impl:         Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
+
+    Returns:
+        Tensor of the same shape and datatype as `x`.
+    """
+
+    k = _setup_kernel(k) * gain
+    p = k.shape[0] - 1
+    return _simple_upfirdn_2d(x, k, pad0=(p+1)//2, pad1=p//2, data_format=data_format, impl=impl)
+
+#----------------------------------------------------------------------------
+
+def upsample_2d(x, k=None, factor=2, gain=1, data_format='NCHW', impl='cuda'):
+    r"""Upsample a batch of 2D images with the given filter.
+
+    Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]`
+    and upsamples each image with the given filter. The filter is normalized so that
+    if the input pixels are constant, they will be scaled by the specified `gain`.
+    Pixels outside the image are assumed to be zero, and the filter is padded with
+    zeros so that its shape is a multiple of the upsampling factor.
+
+    Args:
+        x:            Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.
+        k:            FIR filter of the shape `[firH, firW]` or `[firN]` (separable).
+                      The default is `[1] * factor`, which corresponds to nearest-neighbor
+                      upsampling.
+        factor:       Integer upsampling factor (default: 2).
+        gain:         Scaling factor for signal magnitude (default: 1.0).
+        data_format:  `'NCHW'` or `'NHWC'` (default: `'NCHW'`).
+        impl:         Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
+
+    Returns:
+        Tensor of the shape `[N, C, H * factor, W * factor]` or
+        `[N, H * factor, W * factor, C]`, and same datatype as `x`.
+    """
+
+    assert isinstance(factor, int) and factor >= 1
+    if k is None:
+        k = [1] * factor
+    k = _setup_kernel(k) * (gain * (factor ** 2))
+    p = k.shape[0] - factor
+    return _simple_upfirdn_2d(x, k, up=factor, pad0=(p+1)//2+factor-1, pad1=p//2, data_format=data_format, impl=impl)
+
+#----------------------------------------------------------------------------
+
+def downsample_2d(x, k=None, factor=2, gain=1, data_format='NCHW', impl='cuda'):
+    r"""Downsample a batch of 2D images with the given filter.
+
+    Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]`
+    and downsamples each image with the given filter. The filter is normalized so that
+    if the input pixels are constant, they will be scaled by the specified `gain`.
+    Pixels outside the image are assumed to be zero, and the filter is padded with
+    zeros so that its shape is a multiple of the downsampling factor.
+
+    Args:
+        x:            Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.
+        k:            FIR filter of the shape `[firH, firW]` or `[firN]` (separable).
+                      The default is `[1] * factor`, which corresponds to average pooling.
+        factor:       Integer downsampling factor (default: 2).
+        gain:         Scaling factor for signal magnitude (default: 1.0).
+        data_format:  `'NCHW'` or `'NHWC'` (default: `'NCHW'`).
+        impl:         Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
+
+    Returns:
+        Tensor of the shape `[N, C, H // factor, W // factor]` or
+        `[N, H // factor, W // factor, C]`, and same datatype as `x`.
+    """
+
+    assert isinstance(factor, int) and factor >= 1
+    if k is None:
+        k = [1] * factor
+    k = _setup_kernel(k) * gain
+    p = k.shape[0] - factor
+    return _simple_upfirdn_2d(x, k, down=factor, pad0=(p+1)//2, pad1=p//2, data_format=data_format, impl=impl)
+
+#----------------------------------------------------------------------------
+
+def upsample_conv_2d(x, w, k=None, factor=2, gain=1, data_format='NCHW', impl='cuda'):
+    r"""Fused `upsample_2d()` followed by `tf.nn.conv2d()`.
+
+    Padding is performed only once at the beginning, not between the operations.
+    The fused op is considerably more efficient than performing the same calculation
+    using standard TensorFlow ops. It supports gradients of arbitrary order.
+
+    Args:
+        x:            Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.
+        w:            Weight tensor of the shape `[filterH, filterW, inChannels, outChannels]`.
+                      Grouped convolution can be performed by `inChannels = x.shape[0] // numGroups`.
+        k:            FIR filter of the shape `[firH, firW]` or `[firN]` (separable).
+                      The default is `[1] * factor`, which corresponds to nearest-neighbor
+                      upsampling.
+        factor:       Integer upsampling factor (default: 2).
+        gain:         Scaling factor for signal magnitude (default: 1.0).
+        data_format:  `'NCHW'` or `'NHWC'` (default: `'NCHW'`).
+        impl:         Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
+
+    Returns:
+        Tensor of the shape `[N, C, H * factor, W * factor]` or
+        `[N, H * factor, W * factor, C]`, and same datatype as `x`.
+    """
+
+    assert isinstance(factor, int) and factor >= 1
+
+    # Check weight shape.
+    w = tf.convert_to_tensor(w)
+    assert w.shape.rank == 4
+    convH = w.shape[0]
+    convW = w.shape[1]
+    inC = _shape(w, 2)
+    outC = _shape(w, 3)
+    assert convW == convH
+
+    # Setup filter kernel.
+    if k is None:
+        k = [1] * factor
+    k = _setup_kernel(k) * (gain * (factor ** 2))
+    p = (k.shape[0] - factor) - (convW - 1)
+
+    # Determine data dimensions.
+    if data_format == 'NCHW':
+        stride = [1, 1, factor, factor]
+        output_shape = [_shape(x, 0), outC, (_shape(x, 2) - 1) * factor + convH, (_shape(x, 3) - 1) * factor + convW]
+        num_groups = _shape(x, 1) // inC
+    else:
+        stride = [1, factor, factor, 1]
+        output_shape = [_shape(x, 0), (_shape(x, 1) - 1) * factor + convH, (_shape(x, 2) - 1) * factor + convW, outC]
+        num_groups = _shape(x, 3) // inC
+
+    # Transpose weights.
+    w = tf.reshape(w, [convH, convW, inC, num_groups, -1])
+    w = tf.transpose(w[::-1, ::-1], [0, 1, 4, 3, 2])
+    w = tf.reshape(w, [convH, convW, -1, num_groups * inC])
+
+    # Execute.
+    x = tf.nn.conv2d_transpose(x, w, output_shape=output_shape, strides=stride, padding='VALID', data_format=data_format)
+    return _simple_upfirdn_2d(x, k, pad0=(p+1)//2+factor-1, pad1=p//2+1, data_format=data_format, impl=impl)
+
+#----------------------------------------------------------------------------
+
+def conv_downsample_2d(x, w, k=None, factor=2, gain=1, data_format='NCHW', impl='cuda'):
+    r"""Fused `tf.nn.conv2d()` followed by `downsample_2d()`.
+
+    Padding is performed only once at the beginning, not between the operations.
+    The fused op is considerably more efficient than performing the same calculation
+    using standard TensorFlow ops. It supports gradients of arbitrary order.
+
+    Args:
+        x:            Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.
+        w:            Weight tensor of the shape `[filterH, filterW, inChannels, outChannels]`.
+                      Grouped convolution can be performed by `inChannels = x.shape[0] // numGroups`.
+        k:            FIR filter of the shape `[firH, firW]` or `[firN]` (separable).
+                      The default is `[1] * factor`, which corresponds to average pooling.
+        factor:       Integer downsampling factor (default: 2).
+        gain:         Scaling factor for signal magnitude (default: 1.0).
+        data_format:  `'NCHW'` or `'NHWC'` (default: `'NCHW'`).
+        impl:         Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
+
+    Returns:
+        Tensor of the shape `[N, C, H // factor, W // factor]` or
+        `[N, H // factor, W // factor, C]`, and same datatype as `x`.
+    """
+
+    assert isinstance(factor, int) and factor >= 1
+    w = tf.convert_to_tensor(w)
+    convH, convW, _inC, _outC = w.shape.as_list()
+    assert convW == convH
+    if k is None:
+        k = [1] * factor
+    k = _setup_kernel(k) * gain
+    p = (k.shape[0] - factor) + (convW - 1)
+    if data_format == 'NCHW':
+        s = [1, 1, factor, factor]
+    else:
+        s = [1, factor, factor, 1]
+    x = _simple_upfirdn_2d(x, k, pad0=(p+1)//2, pad1=p//2, data_format=data_format, impl=impl)
+    return tf.nn.conv2d(x, w, strides=s, padding='VALID', data_format=data_format)
+
+#----------------------------------------------------------------------------
+# Internal helper funcs.
+
+def _shape(tf_expr, dim_idx):
+    if tf_expr.shape.rank is not None:
+        dim = tf_expr.shape[dim_idx]
+        if dim is not None:
+            return dim
+    return tf.shape(tf_expr)[dim_idx]
+
+def _setup_kernel(k):
+    k = np.asarray(k, dtype=np.float32)
+    if k.ndim == 1:
+        k = np.outer(k, k)
+    k /= np.sum(k)
+    assert k.ndim == 2
+    assert k.shape[0] == k.shape[1]
+    return k
+
+def _simple_upfirdn_2d(x, k, up=1, down=1, pad0=0, pad1=0, data_format='NCHW', impl='cuda'):
+    assert data_format in ['NCHW', 'NHWC']
+    assert x.shape.rank == 4
+    y = x
+    if data_format == 'NCHW':
+        y = tf.reshape(y, [-1, _shape(y, 2), _shape(y, 3), 1])
+    y = upfirdn_2d(y, k, upx=up, upy=up, downx=down, downy=down, padx0=pad0, padx1=pad1, pady0=pad0, pady1=pad1, impl=impl)
+    if data_format == 'NCHW':
+        y = tf.reshape(y, [-1, _shape(x, 1), _shape(y, 1), _shape(y, 2)])
+    return y
+
+#----------------------------------------------------------------------------

stylegan2/layers/cuda/upfirdn_2d_v2.py

@@ -0,0 +1,146 @@
+import os
+import numpy as np
+import tensorflow as tf
+from stylegan2.layers.cuda import custom_ops
+
+
+def _get_plugin():
+    loc = os.path.dirname(os.path.abspath(__file__))
+    cu_fn = 'upfirdn_2d.cu'
+    return custom_ops.get_plugin(os.path.join(loc, cu_fn))
+
+
+def _setup_kernel(k):
+    k = np.asarray(k, dtype=np.float32)
+    if k.ndim == 1:
+        k = np.outer(k, k)
+    k /= np.sum(k)
+    assert k.ndim == 2
+    assert k.shape[0] == k.shape[1]
+    return k
+
+
+def compute_paddings(resample_kernel, convW, up, down, is_conv, factor=2, gain=1):
+    assert not (up and down)
+
+    k = [1] * factor if resample_kernel is None else resample_kernel
+    if up:
+        k = _setup_kernel(k) * (gain * (factor ** 2))
+        if is_conv:
+            p = (k.shape[0] - factor) - (convW - 1)
+            pad0 = (p + 1) // 2 + factor - 1
+            pad1 = p // 2 + 1
+        else:
+            p = k.shape[0] - factor
+            pad0 = (p + 1) // 2 + factor - 1
+            pad1 = p // 2
+    elif down:
+        k = _setup_kernel(k) * gain
+        if is_conv:
+            p = (k.shape[0] - factor) + (convW - 1)
+            pad0 = (p + 1) // 2
+            pad1 = p // 2
+        else:
+            p = k.shape[0] - factor
+            pad0 = (p + 1) // 2
+            pad1 = p // 2
+    else:
+        k = resample_kernel
+        pad0, pad1 = 0, 0
+    return k, pad0, pad1
+
+
+def upsample_2d(x, x_res, pad0, pad1, k, factor=2):
+    assert isinstance(factor, int) and factor >= 1
+    return _simple_upfirdn_2d(x, x_res, k, up=factor, pad0=pad0, pad1=pad1)
+
+
+def downsample_2d(x, x_res, pad0, pad1, k, factor=2):
+    assert isinstance(factor, int) and factor >= 1
+    return _simple_upfirdn_2d(x, x_res, k, down=factor, pad0=pad0, pad1=pad1)
+
+
+def upsample_conv_2d(x, x_res, w, convH, convW, pad0, pad1, k, factor=2):
+    assert isinstance(factor, int) and factor >= 1
+
+    # Check weight shape.
+    w = tf.convert_to_tensor(w)
+    assert w.shape.rank == 4
+    # convH = w.shape[0]
+    # convW = w.shape[1]
+    inC = tf.shape(w)[2]
+    outC = tf.shape(w)[3]
+    assert convW == convH
+
+    # Determine data dimensions.
+    stride = [1, 1, factor, factor]
+    output_shape = [tf.shape(x)[0], outC, (x_res - 1) * factor + convH, (int(x_res/2) - 1) * factor + convW]
+    num_groups = tf.shape(x)[1] // inC
+
+    # Transpose weights.
+    w = tf.reshape(w, [convH, convW, inC, num_groups, -1])
+    w = tf.transpose(w[::-1, ::-1], [0, 1, 4, 3, 2])
+    w = tf.reshape(w, [convH, convW, -1, num_groups * inC])
+
+    # Execute.
+    x = tf.nn.conv2d_transpose(x, w, output_shape=output_shape, strides=stride, padding='VALID', data_format='NCHW')
+    new_x_res = output_shape[2]
+    return _simple_upfirdn_2d(x, new_x_res, k, pad0=pad0, pad1=pad1)
+
+
+def conv_downsample_2d(x, x_res, w, convH, convW, pad0, pad1, k, factor=2):
+    assert isinstance(factor, int) and factor >= 1
+    w = tf.convert_to_tensor(w)
+    # convH, convW, _inC, _outC = w.shape.as_list()
+    assert convW == convH
+
+    s = [1, 1, factor, factor]
+    x = _simple_upfirdn_2d(x, x_res, k, pad0=pad0, pad1=pad1)
+    return tf.nn.conv2d(x, w, strides=s, padding='VALID', data_format='NCHW')
+
+
+def _simple_upfirdn_2d(x, x_res, k, up=1, down=1, pad0=0, pad1=0):
+    assert x.shape.rank == 4
+    y = x
+    y = tf.reshape(y, [-1, tf.shape(x)[2], tf.shape(x)[3], 1])
+    y = upfirdn_2d_cuda(y, k, upx=up, upy=up, downx=down, downy=down, padx0=pad0, padx1=pad1, pady0=pad0, pady1=pad1)
+    y = tf.reshape(y, [-1, tf.shape(x)[1], tf.shape(y)[1], tf.shape(y)[2]])
+    return y
+
+
+def upfirdn_2d_cuda(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1):
+    """Fast CUDA implementation of `upfirdn_2d()` using custom ops."""
+
+    x = tf.convert_to_tensor(x)
+    k = np.asarray(k, dtype=np.float32)
+    majorDim, inH, inW, minorDim = x.shape.as_list()
+    kernelH, kernelW = k.shape
+    assert inW >= 1 and inH >= 1
+    assert kernelW >= 1 and kernelH >= 1
+    assert isinstance(upx, int) and isinstance(upy, int)
+    assert isinstance(downx, int) and isinstance(downy, int)
+    assert isinstance(padx0, int) and isinstance(padx1, int)
+    assert isinstance(pady0, int) and isinstance(pady1, int)
+
+    outW = (inW * upx + padx0 + padx1 - kernelW) // downx + 1
+    outH = (inH * upy + pady0 + pady1 - kernelH) // downy + 1
+    assert outW >= 1 and outH >= 1
+
+    kc = tf.constant(k, dtype=x.dtype)
+    gkc = tf.constant(k[::-1, ::-1], dtype=x.dtype)
+    gpadx0 = kernelW - padx0 - 1
+    gpady0 = kernelH - pady0 - 1
+    gpadx1 = inW * upx - outW * downx + padx0 - upx + 1
+    gpady1 = inH * upy - outH * downy + pady0 - upy + 1
+
+    @tf.custom_gradient
+    def func(x):
+        y = _get_plugin().up_fir_dn2d(x=x, k=kc, upx=upx, upy=upy, downx=downx, downy=downy, padx0=padx0, padx1=padx1, pady0=pady0, pady1=pady1)
+        y.set_shape([majorDim, outH, outW, minorDim])
+        @tf.custom_gradient
+        def grad(dy):
+            dx = _get_plugin().up_fir_dn2d(x=dy, k=gkc, upx=downx, upy=downy, downx=upx, downy=upy, padx0=gpadx0, padx1=gpadx1, pady0=gpady0, pady1=gpady1)
+            dx.set_shape([majorDim, inH, inW, minorDim])
+            return dx, func
+        return y, grad
+    return func(x)

stylegan2/layers/dense.py

@@ -0,0 +1,37 @@
+import tensorflow as tf
+
+from stylegan2.layers.commons import compute_runtime_coef
+
+
+class Dense(tf.keras.layers.Layer):
+    def __init__(self, fmaps, gain, lrmul, **kwargs):
+        super(Dense, self).__init__(**kwargs)
+        self.fmaps = fmaps
+        self.gain = gain
+        self.lrmul = lrmul
+
+    def build(self, input_shape):
+        fan_in = tf.reduce_prod(input_shape[1:])
+        weight_shape = [fan_in, self.fmaps]
+        init_std, self.runtime_coef = compute_runtime_coef(weight_shape, self.gain, self.lrmul)
+
+        w_init = tf.random.normal(shape=weight_shape, mean=0.0, stddev=init_std)
+        self.w = tf.Variable(w_init, name='w', trainable=True)
+
+    def call(self, inputs, training=None, mask=None):
+        weight = self.runtime_coef * self.w
+
+        c = tf.reduce_prod(tf.shape(inputs)[1:])
+        x = tf.reshape(inputs, shape=[-1, c])
+        x = tf.matmul(x, weight)
+        return x
+
+    def get_config(self):
+        config = super(Dense, self).get_config()
+        config.update({
+            'fmaps': self.fmaps,
+            'gain': self.gain,
+            'lrmul': self.lrmul,
+            'runtime_coef': self.runtime_coef,
+        })
+        return config

stylegan2/layers/embedding.py

@@ -0,0 +1,24 @@
+import tensorflow as tf
+
+
+class LabelEmbedding(tf.keras.layers.Layer):
+    def __init__(self, embed_dim, **kwargs):
+        super(LabelEmbedding, self).__init__(**kwargs)
+        self.embed_dim = embed_dim
+
+    def build(self, input_shape):
+        weight_shape = [input_shape[1], self.embed_dim]
+        # tf 1.15 mean(0.0), std(1.0) default value of tf.initializers.random_normal()
+        w_init = tf.random.normal(shape=weight_shape, mean=0.0, stddev=1.0)
+        self.w = tf.Variable(w_init, name='w', trainable=True)
+
+    def call(self, inputs, training=None, mask=None):
+        x = tf.matmul(inputs, self.w)
+        return x
+
+    def get_config(self):
+        config = super(LabelEmbedding, self).get_config()
+        config.update({
+            'embed_dim': self.embed_dim,
+        })
+        return config

stylegan2/layers/from_rgb.py

@@ -0,0 +1,28 @@
+import tensorflow as tf
+
+from stylegan2.layers.conv import Conv2D
+from stylegan2.layers.bias_act import BiasAct
+
+
+class FromRGB(tf.keras.layers.Layer):
+    def __init__(self, fmaps, res, **kwargs):
+        super(FromRGB, self).__init__(**kwargs)
+        self.fmaps = fmaps
+        self.res = res
+
+        self.conv = Conv2D(in_res=res, in_fmaps=3, fmaps=self.fmaps, kernel=1, up=False, down=False,
+                           resample_kernel=None, gain=1.0, lrmul=1.0, name='conv')
+        self.apply_bias_act = BiasAct(lrmul=1.0, act='lrelu', name='bias')
+
+    def call(self, inputs, training=None, mask=None):
+        y = self.conv(inputs)
+        y = self.apply_bias_act(y)
+        return y
+
+    def get_config(self):
+        config = super(FromRGB, self).get_config()
+        config.update({
+            'fmaps': self.fmaps,
+            'res': self.res,
+        })
+        return config

stylegan2/layers/mapping_block.py

@@ -0,0 +1,58 @@
+import tensorflow as tf
+
+from stylegan2.layers.embedding import LabelEmbedding
+from stylegan2.layers.dense import Dense
+from stylegan2.layers.bias_act import BiasAct
+
+
+class Mapping(tf.keras.layers.Layer):
+    def __init__(self, w_dim, labels_dim, n_mapping, name, **kwargs):
+        super(Mapping, self).__init__(name=name, **kwargs)
+        self.w_dim = w_dim
+        self.labels_dim = labels_dim
+        self.n_mapping = n_mapping
+        self.gain = 1.0
+        self.lrmul = 0.01
+
+        if self.labels_dim != self.w_dim:
+            self.labels_embedding = LabelEmbedding(embed_dim=self.w_dim, name='labels_embedding')
+        else:
+            self.labels_embedding = tf.keras.layers.Lambda(lambda x: x)
+
+        self.normalize = tf.keras.layers.Lambda(lambda x: x * tf.math.rsqrt(tf.reduce_mean(tf.square(x), axis=1, keepdims=True) + 1e-8))
+
+        self.dense_layers = list()
+        self.bias_act_layers = list()
+        for ii in range(self.n_mapping):
+            self.dense_layers.append(Dense(w_dim, gain=self.gain, lrmul=self.lrmul, name='dense_{:d}'.format(ii)))
+            self.bias_act_layers.append(BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_{:d}'.format(ii)))
+
+    def call(self, inputs, training=None, mask=None):
+        latents, labels = inputs
+        x = latents
+
+        # embed label if any
+        if self.labels_dim > 0:
+            y = self.labels_embedding(labels)
+            x = tf.concat([x, y], axis=1)
+
+        # normalize inputs
+        x = self.normalize(x)
+
+        # apply mapping blocks
+        for dense, apply_bias_act in zip(self.dense_layers, self.bias_act_layers):
+            x = dense(x)
+            x = apply_bias_act(x)
+
+        return x
+
+    def get_config(self):
+        config = super(Mapping, self).get_config()
+        config.update({
+            'w_dim': self.w_dim,
+            'labels_dim': self.labels_dim,
+            'n_mapping': self.n_mapping,
+            'gain': self.gain,
+            'lrmul': self.lrmul,
+        })
+        return config

stylegan2/layers/mini_batch_std.py

@@ -0,0 +1,33 @@
+import tensorflow as tf
+
+
+class MinibatchStd(tf.keras.layers.Layer):
+    def __init__(self, group_size, num_new_features, **kwargs):
+        super(MinibatchStd, self).__init__(**kwargs)
+        self.group_size = group_size
+        self.num_new_features = num_new_features
+
+    def call(self, inputs, training=None, mask=None):
+        s = tf.shape(inputs)
+        group_size = tf.minimum(self.group_size, s[0])
+
+        y = tf.reshape(inputs, [group_size, -1, self.num_new_features, s[1] // self.num_new_features, s[2], s[3]])
+        y = tf.cast(y, tf.float32)
+        y -= tf.reduce_mean(y, axis=0, keepdims=True)
+        y = tf.reduce_mean(tf.square(y), axis=0)
+        y = tf.sqrt(y + 1e-8)
+        y = tf.reduce_mean(y, axis=[2, 3, 4], keepdims=True)
+        y = tf.reduce_mean(y, axis=[2])
+        y = tf.cast(y, inputs.dtype)
+        y = tf.tile(y, [group_size, 1, s[2], s[3]])
+
+        x = tf.concat([inputs, y], axis=1)
+        return x
+
+    def get_config(self):
+        config = super(MinibatchStd, self).get_config()
+        config.update({
+            'group_size': self.group_size,
+            'num_new_features': self.num_new_features,
+        })
+        return config

stylegan2/layers/modulated_conv2d.py

@@ -0,0 +1,105 @@
+import tensorflow as tf
+
+from stylegan2.layers.commons import compute_runtime_coef
+from stylegan2.layers.dense import Dense
+from stylegan2.layers.bias_act import BiasAct
+from stylegan2.layers.cuda.upfirdn_2d_v2 import upsample_conv_2d, conv_downsample_2d, compute_paddings
+
+
+class ModulatedConv2D(tf.keras.layers.Layer):
+    def __init__(self, in_res, in_fmaps, fmaps, kernel, up, down, demodulate, resample_kernel, gain, lrmul,
+                 fused_modconv, **kwargs):
+        super(ModulatedConv2D, self).__init__(**kwargs)
+        assert not (up and down)
+
+        self.in_res = in_res
+        self.in_fmaps = in_fmaps
+        self.fmaps = fmaps
+        self.kernel = kernel
+        self.demodulate = demodulate
+        self.up = up
+        self.down = down
+        self.fused_modconv = fused_modconv
+        self.gain = gain
+        self.lrmul = lrmul
+        # self.resample_kernel = resample_kernel
+
+        self.k, self.pad0, self.pad1 = compute_paddings(resample_kernel, self.kernel, up, down, is_conv=True)
+
+        # self.factor = 2
+        self.mod_dense = Dense(self.in_fmaps, gain=1.0, lrmul=1.0, name='mod_dense')
+        self.mod_bias = BiasAct(lrmul=1.0, act='linear', name='mod_bias')
+
+    def build(self, input_shape):
+        # x_shape, w_shape = input_shape[0], input_shape[1]
+        weight_shape = [self.kernel, self.kernel, self.in_fmaps, self.fmaps]
+        init_std, self.runtime_coef = compute_runtime_coef(weight_shape, self.gain, self.lrmul)
+
+        # [kkIO]
+        w_init = tf.random.normal(shape=weight_shape, mean=0.0, stddev=init_std)
+        self.w = tf.Variable(w_init, name='w', trainable=True)
+
+    def call(self, inputs, training=None, mask=None):
+        x, y = inputs
+        # height, width = tf.shape(x)[2], tf.shape(x)[3]
+
+        # prepare weights: [BkkIO] Introduce minibatch dimension
+        w = self.runtime_coef * self.w
+        ww = w[tf.newaxis]
+
+        # Modulate
+        s = self.mod_dense(y)           # [BI]
+        s = self.mod_bias(s) + 1.0      # [BI]
+        ww *= s[:, tf.newaxis, tf.newaxis, :, tf.newaxis]  # [BkkIO]
+
+        if self.demodulate:
+            d = tf.math.rsqrt(tf.reduce_sum(tf.square(ww), axis=[1, 2, 3]) + 1e-8)  # [BO]
+            ww *= d[:, tf.newaxis, tf.newaxis, tf.newaxis, :]                       # [BkkIO]
+
+        if self.fused_modconv:
+            # Fused => reshape minibatch to convolution groups
+            x_shape = tf.shape(x)
+            ww_shape = tf.shape(ww)
+            x = tf.reshape(x, [1, -1, x_shape[2], x_shape[3]])
+            w = tf.reshape(tf.transpose(ww, [1, 2, 3, 0, 4]), [ww_shape[1], ww_shape[2], ww_shape[3], -1])
+        else:
+            # [BIhw] Not fused => scale input activations
+            x *= s[:, :, tf.newaxis, tf.newaxis]
+
+        # Convolution with optional up/downsampling.
+        if self.up:
+            x = upsample_conv_2d(x, self.in_res, w, self.kernel, self.kernel, self.pad0, self.pad1, self.k)
+        elif self.down:
+            x = conv_downsample_2d(x, self.in_res, w, self.kernel, self.kernel, self.pad0, self.pad1, self.k)
+        else:
+            x = tf.nn.conv2d(x, w, data_format='NCHW', strides=[1, 1, 1, 1], padding='SAME')
+
+        # Reshape/scale output
+        if self.fused_modconv:
+            # Fused => reshape convolution groups back to minibatch
+            x_shape = tf.shape(x)
+            x = tf.reshape(x, [-1, self.fmaps, x_shape[2], x_shape[3]])
+        elif self.demodulate:
+            # [BOhw] Not fused => scale output activations
+            x *= d[:, :, tf.newaxis, tf.newaxis]
+        return x
+
+    def get_config(self):
+        config = super(ModulatedConv2D, self).get_config()
+        config.update({
+            'in_res': self.in_res,
+            'in_fmaps': self.in_fmaps,
+            'fmaps': self.fmaps,
+            'kernel': self.kernel,
+            'demodulate': self.demodulate,
+            'fused_modconv': self.fused_modconv,
+            'gain': self.gain,
+            'lrmul': self.lrmul,
+            'up': self.up,
+            'down': self.down,
+            'k': self.k,
+            'pad0': self.pad0,
+            'pad1': self.pad1,
+            'runtime_coef': self.runtime_coef,
+        })
+        return config

stylegan2/layers/noise.py

@@ -0,0 +1,25 @@
+import tensorflow as tf
+
+
+class Noise(tf.keras.layers.Layer):
+    def __init__(self, **kwargs):
+        super(Noise, self).__init__(**kwargs)
+
+    def build(self, input_shape):
+        w_init = tf.zeros(shape=(), dtype=tf.dtypes.float32)
+        self.noise_strength = tf.Variable(w_init, trainable=True, name='w')
+
+    def call(self, inputs, noise=None, training=None, mask=None):
+        x_shape = tf.shape(inputs)
+
+        # noise: [1, 1, x_shape[2], x_shape[3]] or None
+        if noise is None:
+            noise = tf.random.normal(shape=(x_shape[0], 1, x_shape[2], x_shape[3]), dtype=tf.dtypes.float32)
+
+        x = inputs + noise * self.noise_strength
+        return x
+
+    def get_config(self):
+        config = super(Noise, self).get_config()
+        config.update({})
+        return config

stylegan2/layers/synthesis_block.py

@@ -0,0 +1,162 @@
+import tensorflow as tf
+
+from stylegan2.layers.modulated_conv2d import ModulatedConv2D
+from stylegan2.layers.bias_act import BiasAct
+from stylegan2.layers.noise import Noise
+from stylegan2.layers.to_rgb import ToRGB
+from stylegan2.layers.cuda.upfirdn_2d_v2 import upsample_2d, compute_paddings
+
+
+class SynthesisConstBlock(tf.keras.layers.Layer):
+    def __init__(self, fmaps, res, **kwargs):
+        super(SynthesisConstBlock, self).__init__(**kwargs)
+        assert res == 4
+        self.res = res
+        self.fmaps = fmaps
+        self.gain = 1.0
+        self.lrmul = 1.0
+
+        # conv block
+        self.conv = ModulatedConv2D(in_res=res, in_fmaps=self.fmaps, fmaps=self.fmaps, kernel=3, up=False, down=False,
+                                    demodulate=True, resample_kernel=[1, 3, 3, 1], gain=self.gain, lrmul=self.lrmul,
+                                    fused_modconv=True, name='conv')
+        self.apply_noise = Noise(name='noise')
+        self.apply_bias_act = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias')
+
+    def build(self, input_shape):
+        # starting const variable
+        # tf 1.15 mean(0.0), std(1.0) default value of tf.initializers.random_normal()
+        const_init = tf.random.normal(shape=(1, self.fmaps, self.res, int(self.res/2)), mean=0.0, stddev=1.0)
+        self.const = tf.Variable(const_init, name='const', trainable=True)
+
+    def call(self, inputs, training=None, mask=None):
+        w0 = inputs
+        batch_size = tf.shape(w0)[0]
+
+        # const block
+        x = tf.tile(self.const, [batch_size, 1, 1, 1])
+
+        # conv block
+        x = self.conv([x, w0])
+        x = self.apply_noise(x)
+        x = self.apply_bias_act(x)
+        return x
+
+    def get_config(self):
+        config = super(SynthesisConstBlock, self).get_config()
+        config.update({
+            'res': self.res,
+            'fmaps': self.fmaps,
+            'gain': self.gain,
+            'lrmul': self.lrmul,
+        })
+        return config
+
+
+class SynthesisBlock(tf.keras.layers.Layer):
+    def __init__(self, in_ch, fmaps, res, **kwargs):
+        super(SynthesisBlock, self).__init__(**kwargs)
+        self.in_ch = in_ch
+        self.fmaps = fmaps
+        self.res = res
+        self.gain = 1.0
+        self.lrmul = 1.0
+
+        # conv0 up
+        self.conv_0 = ModulatedConv2D(in_res=res//2, in_fmaps=self.in_ch, fmaps=self.fmaps, kernel=3, up=True, down=False,
+                                      demodulate=True, resample_kernel=[1, 3, 3, 1], gain=self.gain, lrmul=self.lrmul,
+                                      fused_modconv=True, name='conv_0')
+        self.apply_noise_0 = Noise(name='noise_0')
+        self.apply_bias_act_0 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_0')
+
+        # conv block
+        self.conv_1 = ModulatedConv2D(in_res=res, in_fmaps=self.fmaps, fmaps=self.fmaps, kernel=3, up=False, down=False,
+                                      demodulate=True, resample_kernel=[1, 3, 3, 1], gain=self.gain, lrmul=self.lrmul,
+                                      fused_modconv=True, name='conv_1')
+        self.apply_noise_1 = Noise(name='noise_1')
+        self.apply_bias_act_1 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_1')
+
+    def call(self, inputs, training=None, mask=None):
+        x, w0, w1 = inputs
+
+        # conv0 up
+        x = self.conv_0([x, w0])
+        x = self.apply_noise_0(x)
+        x = self.apply_bias_act_0(x)
+
+        # conv block
+        x = self.conv_1([x, w1])
+        x = self.apply_noise_1(x)
+        x = self.apply_bias_act_1(x)
+        return x
+
+    def get_config(self):
+        config = super(SynthesisBlock, self).get_config()
+        config.update({
+            'in_ch': self.in_ch,
+            'res': self.res,
+            'fmaps': self.fmaps,
+            'gain': self.gain,
+            'lrmul': self.lrmul,
+        })
+        return config
+
+
+class Synthesis(tf.keras.layers.Layer):
+    def __init__(self, resolutions, featuremaps, name, **kwargs):
+        super(Synthesis, self).__init__(name=name, **kwargs)
+        self.resolutions = resolutions
+        self.featuremaps = featuremaps
+        # self.resample_kernel = [1, 3, 3, 1]
+
+        self.k, self.pad0, self.pad1 = compute_paddings([1, 3, 3, 1], None, up=True, down=False, is_conv=False)
+
+        # initial layer
+        res, n_f = resolutions[0], featuremaps[0]
+        self.initial_block = SynthesisConstBlock(fmaps=n_f, res=res, name='{:d}x{:d}/const'.format(res, res))
+        self.initial_torgb = ToRGB(in_ch=n_f, res=res, name='{:d}x{:d}/ToRGB'.format(res, res))
+
+        # stack generator block with lerp block
+        prev_n_f = n_f
+        self.blocks = list()
+        self.torgbs = list()
+        for res, n_f in zip(self.resolutions[1:], self.featuremaps[1:]):
+            self.blocks.append(SynthesisBlock(in_ch=prev_n_f, fmaps=n_f, res=res,
+                                              name='{:d}x{:d}/block'.format(res, res)))
+            self.torgbs.append(ToRGB(in_ch=n_f, res=res * 2, name='{:d}x{:d}/ToRGB'.format(res, res)))
+            prev_n_f = n_f
+
+    def call(self, inputs, training=None, mask=None):
+        w_broadcasted = inputs
+
+        # initial layer
+        w0, w1 = w_broadcasted[:, 0], w_broadcasted[:, 1]
+        x = self.initial_block(w0)
+        y = self.initial_torgb([x, w1])
+
+        layer_index = 1
+        for block, torgb in zip(self.blocks, self.torgbs):
+            w0 = w_broadcasted[:, layer_index]
+            w1 = w_broadcasted[:, layer_index + 1]
+            w2 = w_broadcasted[:, layer_index + 2]
+
+            y_res = block.res // 2
+            x = block([x, w0, w1])
+            y = upsample_2d(y, y_res, self.pad0, self.pad1, self.k)
+            y = y + torgb([x, w2])
+
+            layer_index += 2
+
+        images_out = y
+        return images_out
+
+    def get_config(self):
+        config = super(Synthesis, self).get_config()
+        config.update({
+            'resolutions': self.resolutions,
+            'featuremaps': self.featuremaps,
+            'k': self.k,
+            'pad0': self.pad0,
+            'pad1': self.pad1,
+        })
+        return config

stylegan2/layers/to_rgb.py

@@ -0,0 +1,30 @@
+import tensorflow as tf
+
+from stylegan2.layers.modulated_conv2d import ModulatedConv2D
+from stylegan2.layers.bias_act import BiasAct
+
+
+class ToRGB(tf.keras.layers.Layer):
+    def __init__(self, in_ch, res, **kwargs):
+        super(ToRGB, self).__init__(**kwargs)
+        self.in_ch = in_ch
+        self.res = res
+
+        self.conv = ModulatedConv2D(in_res=res, in_fmaps=in_ch, fmaps=3, kernel=1, up=False, down=False, demodulate=False,
+                                    resample_kernel=None, gain=1.0, lrmul=1.0, fused_modconv=True, name='conv')
+        self.apply_bias = BiasAct(lrmul=1.0, act='linear', name='bias')
+
+    def call(self, inputs, training=None, mask=None):
+        x, w = inputs
+
+        x = self.conv([x, w])
+        x = self.apply_bias(x)
+        return x
+
+    def get_config(self):
+        config = super(ToRGB, self).get_config()
+        config.update({
+            'in_ch': self.in_ch,
+            'res': self.res,
+        })
+        return config

stylegan2/utils.py

@@ -0,0 +1,112 @@
+from typing import Any
+import numpy as np
+import tensorflow as tf
+
+
+class EasyDict(dict):
+    """Convenience class that behaves like a dict but allows access with the attribute syntax."""
+
+    def __getattr__(self, name: str) -> Any:
+        try:
+            return self[name]
+        except KeyError:
+            raise AttributeError(name)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        self[name] = value
+
+    def __delattr__(self, name: str) -> None:
+        del self[name]
+
+
+def lerp(a, b, t):
+    out = a + (b - a) * t
+    return out
+
+
+def lerp_clip(a, b, t):
+    out = a + (b - a) * tf.clip_by_value(t, 0.0, 1.0)
+    return out
+
+
+def adjust_dynamic_range(images, range_in, range_out, out_dtype):
+    scale = (range_out[1] - range_out[0]) / (range_in[1] - range_in[0])
+    bias = range_out[0] - range_in[0] * scale
+    images = images * scale + bias
+    images = tf.clip_by_value(images, range_out[0], range_out[1])
+    images = tf.cast(images, dtype=out_dtype)
+    return images
+
+
+def random_flip_left_right_nchw(images):
+    s = tf.shape(images)
+    mask = tf.random.uniform([s[0], 1, 1, 1], 0.0, 1.0)
+    mask = tf.tile(mask, [1, s[1], s[2], s[3]])
+    images = tf.where(mask < 0.5, images, tf.reverse(images, axis=[3]))
+    return images
+
+
+def preprocess_fit_train_image(images, res):
+    images = adjust_dynamic_range(images, range_in=(0.0, 255.0), range_out=(-1.0, 1.0), out_dtype=tf.dtypes.float32)
+    images = random_flip_left_right_nchw(images)
+    images.set_shape([None, 3, res, res])
+    return images
+
+
+def postprocess_images(images):
+    images = adjust_dynamic_range(images, range_in=(-1.0, 1.0), range_out=(0.0, 255.0), out_dtype=tf.dtypes.float32)
+    images = tf.transpose(images, [0, 2, 3, 1])
+    images = tf.cast(images, dtype=tf.dtypes.uint8)
+    return images
+
+
+def merge_batch_images(images, res, rows, cols):
+    batch_size = images.shape[0]
+    assert rows * cols == batch_size
+    canvas = np.zeros(shape=[res * rows, res * cols, 3], dtype=np.uint8)
+    for row in range(rows):
+        y_start = row * res
+        for col in range(cols):
+            x_start = col * res
+            index = col + row * cols
+            canvas[y_start:y_start + res, x_start:x_start + res, :] = images[index, :, :, :]
+    return canvas
+
+
+def main():
+    # out_dtype = tf.dtypes.float32
+    # range_in = (0.0, 255.0)
+    # range_out = (-1.0, 1.0)
+    # # images = np.ones((2, 2), dtype=np.float32) * 255.0
+    # images1 = np.ones((2, 2), dtype=np.float32) * 0.0
+    # images2 = np.ones((2, 2), dtype=np.float32) * 127.5
+    # images3 = np.ones((2, 2), dtype=np.float32) * 255.0
+    # adjusted1 = adjust_dynamic_range(images1, range_in, range_out, out_dtype)
+    # adjusted2 = adjust_dynamic_range(images2, range_in, range_out, out_dtype)
+    # adjusted3 = adjust_dynamic_range(images3, range_in, range_out, out_dtype)
+    # print(adjusted1)
+    # print(adjusted2)
+    # print(adjusted3)
+    #
+    # out_dtype = tf.dtypes.uint8
+    # range_in = (-1.0, 1.0)
+    # range_out = (0.0, 255.0)
+    # images1 = np.ones((2, 2), dtype=np.float32) * -1.0
+    # images2 = np.ones((2, 2), dtype=np.float32) * 0.0
+    # images3 = np.ones((2, 2), dtype=np.float32) * 1.0
+    # adjusted1 = adjust_dynamic_range(images1, range_in, range_out, out_dtype)
+    # adjusted2 = adjust_dynamic_range(images2, range_in, range_out, out_dtype)
+    # adjusted3 = adjust_dynamic_range(images3, range_in, range_out, out_dtype)
+    # print(adjusted1)
+    # print(adjusted2)
+    # print(adjusted3)
+
+    batch_size = 8
+    res = 128
+    fake_images = np.ones(shape=(batch_size, res, res, 3), dtype=np.uint8)
+    merge_batch_images(fake_images, res, rows=4, cols=2)
+    return
+
+
+if __name__ == '__main__':
+    main()

stylegan2_ref/custom_layers.py

@@ -0,0 +1,309 @@
+import numpy as np
+import tensorflow as tf
+
+from stylegan2_ref.upfirdn_2d import setup_resample_kernel, upsample_conv_2d, conv_downsample_2d
+
+
+def compute_runtime_coef(weight_shape, gain, lrmul):
+    fan_in = np.prod(weight_shape[:-1])  # [kernel, kernel, fmaps_in, fmaps_out] or [in, out]
+    he_std = gain / np.sqrt(fan_in)
+    init_std = 1.0 / lrmul
+    runtime_coef = he_std * lrmul
+    return init_std, runtime_coef
+
+
+class Dense(tf.keras.layers.Layer):
+    def __init__(self, fmaps, gain, lrmul, **kwargs):
+        super(Dense, self).__init__(**kwargs)
+        self.fmaps = fmaps
+        self.gain = gain
+        self.lrmul = lrmul
+
+    def build(self, input_shape):
+        assert len(input_shape) == 2 or len(input_shape) == 4
+        fan_in = np.prod(input_shape[1:])
+        weight_shape = [fan_in, self.fmaps]
+        init_std, runtime_coef = compute_runtime_coef(weight_shape, self.gain, self.lrmul)
+
+        self.runtime_coef = runtime_coef
+
+        w_init = tf.random.normal(shape=weight_shape, mean=0.0, stddev=init_std)
+        self.w = tf.Variable(w_init, name='w', trainable=True)
+
+    def call(self, inputs, training=None, mask=None):
+        weight = self.runtime_coef * self.w
+        x = tf.keras.layers.Flatten()(inputs)
+        x = tf.matmul(x, weight)
+        return x
+
+    def get_config(self):
+        config = super(Dense, self).get_config()
+        config.update({
+            'fmaps': self.fmaps,
+            'gain': self.gain,
+            'lrmul': self.lrmul,
+        })
+        return config
+
+
+class BiasAct(tf.keras.layers.Layer):
+    def __init__(self, lrmul, act, **kwargs):
+        super(BiasAct, self).__init__(**kwargs)
+        assert act in ['linear', 'lrelu']
+        self.lrmul = lrmul
+
+        if act == 'linear':
+            self.act = tf.keras.layers.Lambda(lambda x: tf.identity(x))
+            self.gain = 1.0
+        else:
+            self.act = tf.keras.layers.LeakyReLU(alpha=0.2)
+            self.gain = np.sqrt(2)
+
+    def build(self, input_shape):
+        self.len2 = True if len(input_shape) == 2 else False
+        b_init = tf.zeros(shape=(input_shape[1],), dtype=tf.dtypes.float32)
+        self.b = tf.Variable(b_init, name='b', trainable=True)
+
+    def call(self, inputs, training=None, mask=None):
+        b = self.lrmul * self.b
+
+        if self.len2:
+            x = inputs + b
+        else:
+            x = inputs + tf.reshape(b, shape=[1, -1, 1, 1])
+        x = self.act(x)
+        x = self.gain * x
+        return x
+
+    def get_config(self):
+        config = super(BiasAct, self).get_config()
+        config.update({
+            'lrmul': self.lrmul,
+            'gain': self.gain,
+            'len2': self.len2,
+        })
+        return config
+
+
+class LeakyReLU(tf.keras.layers.Layer):
+    def __init__(self, **kwargs):
+        super(LeakyReLU, self).__init__(**kwargs)
+        self.alpha = 0.2
+        self.gain = np.sqrt(2)
+
+        self.act = tf.keras.layers.LeakyReLU(alpha=self.alpha)
+
+    def call(self, inputs, training=None, mask=None):
+        x = self.act(inputs)
+        x *= self.gain
+        return x
+
+    def get_config(self):
+        config = super(LeakyReLU, self).get_config()
+        config.update({
+            'alpha': self.alpha,
+            'gain': self.gain,
+        })
+        return config
+
+
+class LabelEmbedding(tf.keras.layers.Layer):
+    def __init__(self, embed_dim, **kwargs):
+        super(LabelEmbedding, self).__init__(**kwargs)
+        self.embed_dim = embed_dim
+
+    def build(self, input_shape):
+        weight_shape = [input_shape[1], self.embed_dim]
+        # tf 1.15 mean(0.0), std(1.0) default value of tf.initializers.random_normal()
+        w_init = tf.random.normal(shape=weight_shape, mean=0.0, stddev=1.0)
+        self.w = tf.Variable(w_init, name='w', trainable=True)
+
+    def call(self, inputs, training=None, mask=None):
+        x = tf.matmul(inputs, self.w)
+        return x
+
+    def get_config(self):
+        config = super(LabelEmbedding, self).get_config()
+        config.update({
+            'embed_dim': self.embed_dim,
+        })
+        return config
+
+
+class Noise(tf.keras.layers.Layer):
+    def __init__(self, **kwargs):
+        super(Noise, self).__init__(**kwargs)
+
+    def build(self, input_shape):
+        self.noise_strength = tf.Variable(initial_value=0.0, dtype=tf.dtypes.float32, trainable=True, name='w')
+
+    def call(self, x, training=None, mask=None):
+        x_shape = tf.shape(x)
+        noise = tf.random.normal(shape=(x_shape[0], 1, x_shape[2], x_shape[3]), dtype=tf.dtypes.float32)
+
+        x += noise * self.noise_strength
+        return x
+
+
+class MinibatchStd(tf.keras.layers.Layer):
+    def __init__(self, group_size, num_new_features, **kwargs):
+        super(MinibatchStd, self).__init__(**kwargs)
+        self.group_size = group_size
+        self.num_new_features = num_new_features
+
+    def call(self, x, training=None, mask=None):
+        group_size = tf.minimum(self.group_size, tf.shape(x)[0])
+        s = tf.shape(x)
+
+        y = tf.reshape(x, [group_size, -1, self.num_new_features, s[1] // self.num_new_features, s[2], s[3]])
+        y = tf.cast(y, tf.float32)
+        y -= tf.reduce_mean(y, axis=0, keepdims=True)
+        y = tf.reduce_mean(tf.square(y), axis=0)
+        y = tf.sqrt(y + 1e-8)
+        y = tf.reduce_mean(y, axis=[2, 3, 4], keepdims=True)
+        y = tf.reduce_mean(y, axis=[2])
+        y = tf.cast(y, x.dtype)
+        y = tf.tile(y, [group_size, 1, s[2], s[3]])
+        return tf.concat([x, y], axis=1)
+
+
+class FusedModConv(tf.keras.layers.Layer):
+    def __init__(self, fmaps, kernel, gain, lrmul, style_fmaps, demodulate, up, down, resample_kernel, **kwargs):
+        super(FusedModConv, self).__init__(**kwargs)
+        assert not (up and down)
+        self.fmaps = fmaps
+        self.kernel = kernel
+        self.gain = gain
+        self.lrmul = lrmul
+        self.style_fmaps = style_fmaps
+        self.demodulate = demodulate
+        self.up = up
+        self.down = down
+        self.factor = 2
+        if resample_kernel is None:
+            resample_kernel = [1] * self.factor
+        self.k = setup_resample_kernel(k=resample_kernel)
+
+        self.mod_dense = Dense(self.style_fmaps, gain=1.0, lrmul=1.0, name='mod_dense')
+        self.mod_bias = BiasAct(lrmul=1.0, act='linear', name='mod_bias')
+
+    def build(self, input_shape):
+        x_shape, w_shape = input_shape[0], input_shape[1]
+        weight_shape = [self.kernel, self.kernel, x_shape[1], self.fmaps]
+        init_std, runtime_coef = compute_runtime_coef(weight_shape, self.gain, self.lrmul)
+
+        self.runtime_coef = runtime_coef
+
+        # [kkIO]
+        w_init = tf.random.normal(shape=weight_shape, mean=0.0, stddev=init_std)
+        self.w = tf.Variable(w_init, name='w', trainable=True)
+
+    def scale_conv_weights(self, w):
+        # convolution kernel weights for fused conv
+        weight = self.runtime_coef * self.w     # [kkIO]
+        weight = weight[np.newaxis]             # [BkkIO]
+
+        # modulation
+        style = self.mod_dense(w)                                   # [BI]
+        style = self.mod_bias(style) + 1.0                          # [BI]
+        weight *= style[:, np.newaxis, np.newaxis, :, np.newaxis]   # [BkkIO]
+
+        # demodulation
+        if self.demodulate:
+            d = tf.math.rsqrt(tf.reduce_sum(tf.square(weight), axis=[1, 2, 3]) + 1e-8)  # [BO]
+            weight *= d[:, np.newaxis, np.newaxis, np.newaxis, :]                       # [BkkIO]
+
+        # weight: reshape, prepare for fused operation
+        new_weight_shape = [tf.shape(weight)[1], tf.shape(weight)[2], tf.shape(weight)[3], -1]      # [kkI(BO)]
+        weight = tf.transpose(weight, [1, 2, 3, 0, 4])                                              # [kkIBO]
+        weight = tf.reshape(weight, shape=new_weight_shape)                                         # [kkI(BO)]
+        return weight
+
+    def call(self, inputs, training=None, mask=None):
+        x, w = inputs
+        height, width = tf.shape(x)[2], tf.shape(x)[3]
+
+        # prepare convolution kernel weights
+        weight = self.scale_conv_weights(w)
+
+        # prepare inputs: reshape minibatch to convolution groups
+        x = tf.reshape(x, [1, -1, height, width])
+
+        if self.up:
+            x = upsample_conv_2d(x, self.k, weight, self.factor, self.gain)
+        elif self.down:
+            x = conv_downsample_2d(x, self.k, weight, self.factor, self.gain)
+        else:
+            x = tf.nn.conv2d(x, weight, data_format='NCHW', strides=[1, 1, 1, 1], padding='SAME')
+
+        # x: reshape back
+        x = tf.reshape(x, [-1, self.fmaps, tf.shape(x)[2], tf.shape(x)[3]])
+        return x
+
+    def get_config(self):
+        config = super(FusedModConv, self).get_config()
+        config.update({
+            'fmaps': self.fmaps,
+            'kernel': self.kernel,
+            'gain': self.gain,
+            'lrmul': self.lrmul,
+            'style_fmaps': self.style_fmaps,
+            'demodulate': self.demodulate,
+            'up': self.up,
+            'down': self.down,
+            'factor': self.factor,
+            'k': self.k,
+        })
+        return config
+
+
+class ResizeConv2D(tf.keras.layers.Layer):
+    def __init__(self, fmaps, kernel, gain, lrmul, up, down, resample_kernel, **kwargs):
+        super(ResizeConv2D, self).__init__(**kwargs)
+        self.fmaps = fmaps
+        self.kernel = kernel
+        self.gain = gain
+        self.lrmul = lrmul
+        self.up = up
+        self.down = down
+        self.factor = 2
+        if resample_kernel is None:
+            resample_kernel = [1] * self.factor
+        self.k = setup_resample_kernel(k=resample_kernel)
+
+    def build(self, input_shape):
+        assert len(input_shape) == 4
+        weight_shape = [self.kernel, self.kernel, input_shape[1], self.fmaps]
+        init_std, runtime_coef = compute_runtime_coef(weight_shape, self.gain, self.lrmul)
+
+        self.runtime_coef = runtime_coef
+
+        # [kernel, kernel, fmaps_in, fmaps_out]
+        w_init = tf.random_normal_initializer(mean=0.0, stddev=init_std)
+        self.w = tf.Variable(initial_value=w_init(shape=weight_shape, dtype='float32'), trainable=True, name='w')
+
+    def call(self, inputs, training=None, mask=None):
+        x = inputs
+        weight = self.runtime_coef * self.w
+
+        if self.up:
+            x = upsample_conv_2d(x, self.k, weight, self.factor, self.gain)
+        elif self.down:
+            x = conv_downsample_2d(x, self.k, weight, self.factor, self.gain)
+        else:
+            x = tf.nn.conv2d(x, weight, data_format='NCHW', strides=[1, 1, 1, 1], padding='SAME')
+        return x
+
+    def get_config(self):
+        config = super(ResizeConv2D, self).get_config()
+        config.update({
+            'fmaps': self.fmaps,
+            'kernel': self.kernel,
+            'gain': self.gain,
+            'lrmul': self.lrmul,
+            'up': self.up,
+            'down': self.down,
+            'factor': self.factor,
+            'k': self.k,
+        })
+        return config

stylegan2_ref/discriminator.py

@@ -0,0 +1,150 @@
+import numpy as np
+import tensorflow as tf
+
+from stylegan2_ref.custom_layers import Dense, ResizeConv2D, BiasAct, LeakyReLU, MinibatchStd
+
+
+class FromRGB(tf.keras.layers.Layer):
+    def __init__(self, fmaps, **kwargs):
+        super(FromRGB, self).__init__(**kwargs)
+        self.fmaps = fmaps
+
+        self.conv = ResizeConv2D(fmaps=self.fmaps, kernel=1, gain=1.0, lrmul=1.0,
+                                 up=False, down=False, resample_kernel=None, name='conv')
+        self.apply_bias_act = BiasAct(lrmul=1.0, act='lrelu', name='bias')
+
+    def call(self, inputs, training=None, mask=None):
+        y = self.conv(inputs)
+        y = self.apply_bias_act(y)
+        return y
+
+
+class DiscriminatorBlock(tf.keras.layers.Layer):
+    def __init__(self, n_f0, n_f1, res, **kwargs):
+        super(DiscriminatorBlock, self).__init__(**kwargs)
+        self.gain = 1.0
+        self.lrmul = 1.0
+        self.n_f0 = n_f0
+        self.n_f1 = n_f1
+        self.res = res
+        self.resnet_scale = 1. / np.sqrt(2.)
+
+        # conv_0
+        self.conv_0 = ResizeConv2D(fmaps=self.n_f0, kernel=3, gain=self.gain, lrmul=self.lrmul,
+                                   up=False, down=False, resample_kernel=None, name='conv_0')
+        self.apply_bias_act_0 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_0')
+
+        # conv_1 down
+        self.conv_1 = ResizeConv2D(fmaps=self.n_f1, kernel=3, gain=self.gain, lrmul=self.lrmul,
+                                   up=False, down=True, resample_kernel=[1, 3, 3, 1], name='conv_1')
+        self.apply_bias_act_1 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_1')
+
+        # resnet skip
+        self.conv_skip = ResizeConv2D(fmaps=self.n_f1, kernel=1, gain=self.gain, lrmul=self.lrmul,
+                                      up=False, down=True, resample_kernel=[1, 3, 3, 1], name='skip')
+
+    def call(self, inputs, training=None, mask=None):
+        x = inputs
+        residual = x
+
+        # conv0
+        x = self.conv_0(x)
+        x = self.apply_bias_act_0(x)
+
+        # conv1 down
+        x = self.conv_1(x)
+        x = self.apply_bias_act_1(x)
+
+        # resnet skip
+        residual = self.conv_skip(residual)
+        x = (x + residual) * self.resnet_scale
+        return x
+
+
+class DiscriminatorLastBlock(tf.keras.layers.Layer):
+    def __init__(self, n_f0, n_f1, res, **kwargs):
+        super(DiscriminatorLastBlock, self).__init__(**kwargs)
+        self.gain = 1.0
+        self.lrmul = 1.0
+        self.n_f0 = n_f0
+        self.n_f1 = n_f1
+        self.res = res
+
+        self.minibatch_std = MinibatchStd(group_size=4, num_new_features=1, name='minibatchstd')
+
+        # conv_0
+        self.conv_0 = ResizeConv2D(fmaps=self.n_f0, kernel=3, gain=self.gain, lrmul=self.lrmul,
+                                   up=False, down=False, resample_kernel=None, name='conv_0')
+        self.apply_bias_act_0 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_0')
+
+        # dense_1
+        self.dense_1 = Dense(self.n_f1, gain=self.gain, lrmul=self.lrmul, name='dense_1')
+        self.apply_bias_act_1 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_1')
+
+    def call(self, x, training=None, mask=None):
+        x = self.minibatch_std(x)
+
+        # conv_0
+        x = self.conv_0(x)
+        x = self.apply_bias_act_0(x)
+
+        # dense_1
+        x = self.dense_1(x)
+        x = self.apply_bias_act_1(x)
+        return x
+
+
+class Discriminator(tf.keras.Model):
+    def __init__(self, d_params, **kwargs):
+        super(Discriminator, self).__init__(**kwargs)
+        # discriminator's (resolutions and featuremaps) are reversed against generator's
+        self.labels_dim = d_params['labels_dim']
+        self.r_resolutions = d_params['resolutions'][::-1]
+        self.r_featuremaps = d_params['featuremaps'][::-1]
+
+        # stack discriminator blocks
+        res0, n_f0 = self.r_resolutions[0], self.r_featuremaps[0]
+        self.initial_fromrgb = FromRGB(fmaps=n_f0, name='{:d}x{:d}/FromRGB'.format(res0, res0))
+        self.blocks = list()
+        for index, (res0, n_f0) in enumerate(zip(self.r_resolutions[:-1], self.r_featuremaps[:-1])):
+            n_f1 = self.r_featuremaps[index + 1]
+            self.blocks.append(DiscriminatorBlock(n_f0=n_f0, n_f1=n_f1, res=res0, name='{:d}x{:d}'.format(res0, res0)))
+
+        # set last discriminator block
+        res = self.r_resolutions[-1]
+        n_f0, n_f1 = self.r_featuremaps[-2], self.r_featuremaps[-1]
+        self.last_block = DiscriminatorLastBlock(n_f0, n_f1, res, name='{:d}x{:d}'.format(res, res))
+
+        # set last dense layer
+        self.last_dense = Dense(1, gain=1.0, lrmul=1.0, name='last_dense')
+        self.last_bias = BiasAct(lrmul=1.0, act='linear', name='last_bias')
+
+        # conditional output
+        self.last_block_condition = DiscriminatorLastBlock(n_f0, n_f1, res, name='{:d}x{:d}'.format(res, res))
+        self.condition_dense = Dense(self.labels_dim, gain=1.0, lrmul=1.0, name='last_dense')
+        self.condition_bias = BiasAct(lrmul=1.0, act='linear', name='last_bias')
+
+    def call(self, inputs, training=None, mask=None):
+        images, _ = inputs
+
+        x = self.initial_fromrgb(images)
+        for block in self.blocks:
+            x = block(x)
+
+        features = x
+
+        # score output
+        x = self.last_block(features)
+        x = self.last_dense(x)
+        x = self.last_bias(x)
+
+        # condition output
+        cond = self.last_block_condition(features)
+        cond = self.condition_dense(cond)
+        cond = self.condition_bias(cond)
+
+        scores_out = x
+        return scores_out, cond
+
+    def compute_output_shape(self, input_shape):
+        return input_shape[0][0], max(self.labels_dim, 1)

stylegan2_ref/generator.py

@@ -0,0 +1,336 @@
+import tensorflow as tf
+
+from stylegan2_ref.utils import lerp
+from stylegan2_ref.upfirdn_2d import setup_resample_kernel, upsample_2d
+from stylegan2_ref.custom_layers import LabelEmbedding, Dense, BiasAct, LeakyReLU, Noise, FusedModConv
+
+
+class ToRGB(tf.keras.layers.Layer):
+    def __init__(self, in_ch, **kwargs):
+        super(ToRGB, self).__init__(**kwargs)
+        self.in_ch = in_ch
+        self.conv = FusedModConv(fmaps=3, kernel=1, gain=1.0, lrmul=1.0, style_fmaps=self.in_ch,
+                                 demodulate=False, up=False, down=False, resample_kernel=None, name='conv')
+        self.apply_bias = BiasAct(lrmul=1.0, act='linear', name='bias')
+
+    def call(self, inputs, training=None, mask=None):
+        x, w = inputs
+        assert x.shape[1] == self.in_ch
+
+        x = self.conv([x, w])
+        x = self.apply_bias(x)
+        return x
+
+    def get_config(self):
+        config = super(ToRGB, self).get_config()
+        config.update({
+            'in_ch': self.in_ch,
+        })
+        return config
+
+
+class Mapping(tf.keras.layers.Layer):
+    def __init__(self, w_dim, labels_dim, n_mapping, **kwargs):
+        super(Mapping, self).__init__(**kwargs)
+        self.w_dim = w_dim
+        self.labels_dim = labels_dim
+        self.n_mapping = n_mapping
+        self.gain = 1.0
+        self.lrmul = 0.01
+
+        if self.labels_dim != self.w_dim:
+            self.labels_embedding = LabelEmbedding(embed_dim=self.w_dim, name='labels_embedding')
+        else:
+            self.labels_embedding = tf.keras.layers.Lambda(lambda x: x)
+
+        self.normalize = tf.keras.layers.Lambda(lambda x: x * tf.math.rsqrt(tf.reduce_mean(tf.square(x), axis=1, keepdims=True) + 1e-8))
+
+        self.dense_layers = list()
+        self.bias_act_layers = list()
+        for ii in range(self.n_mapping):
+            self.dense_layers.append(Dense(w_dim, gain=self.gain, lrmul=self.lrmul, name='dense_{:d}'.format(ii)))
+            self.bias_act_layers.append(BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_{:d}'.format(ii)))
+
+    def call(self, inputs, training=None, mask=None):
+        latents, labels = inputs
+        x = latents
+
+        # embed label if any
+        if self.labels_dim > 0:
+            y = self.labels_embedding(labels)
+            x = tf.concat([x, y], axis=1)
+
+        # normalize inputs
+        x = self.normalize(x)
+
+        # apply mapping blocks
+        for dense, apply_bias_act in zip(self.dense_layers, self.bias_act_layers):
+            x = dense(x)
+            x = apply_bias_act(x)
+
+        return x
+
+    def get_config(self):
+        config = super(Mapping, self).get_config()
+        config.update({
+            'w_dim': self.w_dim,
+            'labels_dim': self.labels_dim,
+            'n_mapping': self.n_mapping,
+            'gain': self.gain,
+            'lrmul': self.lrmul,
+        })
+        return config
+
+
+class SynthesisConstBlock(tf.keras.layers.Layer):
+    def __init__(self, fmaps, res, **kwargs):
+        super(SynthesisConstBlock, self).__init__(**kwargs)
+        assert res == 4
+        self.res = res
+        self.fmaps = fmaps
+        self.gain = 1.0
+        self.lrmul = 1.0
+
+        # conv block
+        self.conv = FusedModConv(fmaps=self.fmaps, kernel=3, gain=self.gain, lrmul=self.lrmul, style_fmaps=self.fmaps,
+                                 demodulate=True, up=False, down=False, resample_kernel=[1, 3, 3, 1], name='conv')
+        self.apply_noise = Noise(name='noise')
+        self.apply_bias_act = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias')
+
+    def build(self, input_shape):
+        # starting const variable
+        # tf 1.15 mean(0.0), std(1.0) default value of tf.initializers.random_normal()
+        const_init = tf.random.normal(shape=(1, self.fmaps, self.res, int(self.res/2)), mean=0.0, stddev=1.0)
+        self.const = tf.Variable(const_init, name='const', trainable=True)
+
+    def call(self, inputs, training=None, mask=None):
+        w0 = inputs
+        batch_size = tf.shape(w0)[0]
+
+        # const block
+        x = tf.tile(self.const, [batch_size, 1, 1, 1])
+
+        # conv block
+        x = self.conv([x, w0])
+        x = self.apply_noise(x)
+        x = self.apply_bias_act(x)
+        return x
+
+    def get_config(self):
+        config = super(SynthesisConstBlock, self).get_config()
+        config.update({
+            'res': self.res,
+            'fmaps': self.fmaps,
+            'gain': self.gain,
+            'lrmul': self.lrmul,
+        })
+        return config
+
+
+class SynthesisBlock(tf.keras.layers.Layer):
+    def __init__(self, in_ch, fmaps, res, **kwargs):
+        super(SynthesisBlock, self).__init__(**kwargs)
+        self.in_ch = in_ch
+        self.fmaps = fmaps
+        self.res = res
+        self.gain = 1.0
+        self.lrmul = 1.0
+
+        # conv0 up
+        self.conv_0 = FusedModConv(fmaps=self.fmaps, kernel=3, gain=self.gain, lrmul=self.lrmul, style_fmaps=self.in_ch,
+                                   demodulate=True, up=True, down=False, resample_kernel=[1, 3, 3, 1], name='conv_0')
+        self.apply_noise_0 = Noise(name='noise_0')
+        self.apply_bias_act_0 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_0')
+
+        # conv block
+        self.conv_1 = FusedModConv(fmaps=self.fmaps, kernel=3, gain=self.gain, lrmul=self.lrmul, style_fmaps=self.fmaps,
+                                   demodulate=True, up=False, down=False, resample_kernel=[1, 3, 3, 1], name='conv_1')
+        self.apply_noise_1 = Noise(name='noise_1')
+        self.apply_bias_act_1 = BiasAct(lrmul=self.lrmul, act='lrelu', name='bias_1')
+
+    def call(self, inputs, training=None, mask=None):
+        x, w0, w1 = inputs
+
+        # conv0 up
+        x = self.conv_0([x, w0])
+        x = self.apply_noise_0(x)
+        x = self.apply_bias_act_0(x)
+
+        # conv block
+        x = self.conv_1([x, w1])
+        x = self.apply_noise_1(x)
+        x = self.apply_bias_act_1(x)
+        return x
+
+    def get_config(self):
+        config = super(SynthesisBlock, self).get_config()
+        config.update({
+            'in_ch': self.in_ch,
+            'res': self.res,
+            'fmaps': self.fmaps,
+            'gain': self.gain,
+            'lrmul': self.lrmul,
+        })
+        return config
+
+
+class Synthesis(tf.keras.layers.Layer):
+    def __init__(self, resolutions, featuremaps, name, **kwargs):
+        super(Synthesis, self).__init__(name=name, **kwargs)
+        self.resolutions = resolutions
+        self.featuremaps = featuremaps
+        self.k = setup_resample_kernel(k=[1, 3, 3, 1])
+
+        # initial layer
+        res, n_f = resolutions[0], featuremaps[0]
+        self.initial_block = SynthesisConstBlock(fmaps=n_f, res=res, name='{:d}x{:d}/const'.format(res, res))
+        self.initial_torgb = ToRGB(in_ch=n_f, name='{:d}x{:d}/ToRGB'.format(res, res))
+
+        # stack generator block with lerp block
+        prev_n_f = n_f
+        self.blocks = list()
+        self.torgbs = list()
+        for res, n_f in zip(self.resolutions[1:], self.featuremaps[1:]):
+            self.blocks.append(SynthesisBlock(in_ch=prev_n_f, fmaps=n_f, res=res,
+                                              name='{:d}x{:d}/block'.format(res, res)))
+            self.torgbs.append(ToRGB(in_ch=n_f, name='{:d}x{:d}/ToRGB'.format(res, res)))
+            prev_n_f = n_f
+
+    def call(self, inputs, training=None, mask=None):
+        w_broadcasted = inputs
+
+        # initial layer
+        w0, w1 = w_broadcasted[:, 0], w_broadcasted[:, 1]
+        x = self.initial_block(w0)
+        y = self.initial_torgb([x, w1])
+
+        layer_index = 1
+        for block, torgb in zip(self.blocks, self.torgbs):
+            w0 = w_broadcasted[:, layer_index]
+            w1 = w_broadcasted[:, layer_index + 1]
+            w2 = w_broadcasted[:, layer_index + 2]
+
+            x = block([x, w0, w1])
+            y = upsample_2d(y, self.k, factor=2, gain=1.0)
+            y = y + torgb([x, w2])
+
+            layer_index += 2
+
+        images_out = y
+        return images_out
+
+    def get_config(self):
+        config = super(Synthesis, self).get_config()
+        config.update({
+            'resolutions': self.resolutions,
+            'featuremaps': self.featuremaps,
+            'k': self.k,
+        })
+        return config
+
+
+class Generator(tf.keras.Model):
+    def __init__(self, g_params, **kwargs):
+        super(Generator, self).__init__(**kwargs)
+
+        self.z_dim = g_params['z_dim']
+        self.w_dim = g_params['w_dim']
+        self.labels_dim = g_params['labels_dim']
+        self.n_mapping = g_params['n_mapping']
+        self.resolutions = g_params['resolutions']
+        self.featuremaps = g_params['featuremaps']
+        self.w_ema_decay = 0.995
+        self.style_mixing_prob = 0.9
+
+        self.n_broadcast = len(self.resolutions) * 2
+        self.mixing_layer_indices = tf.range(self.n_broadcast, dtype=tf.int32)[tf.newaxis, :, tf.newaxis]
+
+        self.g_mapping = Mapping(self.w_dim, self.labels_dim, self.n_mapping, name='g_mapping')
+        self.broadcast = tf.keras.layers.Lambda(lambda x: tf.tile(x[:, tf.newaxis], [1, self.n_broadcast, 1]))
+        self.synthesis = Synthesis(self.resolutions, self.featuremaps, name='g_synthesis')
+
+    def build(self, input_shape):
+        # w_avg
+        self.w_avg = tf.Variable(tf.zeros(shape=[self.w_dim], dtype=tf.dtypes.float32), name='w_avg', trainable=False,
+                                 synchronization=tf.VariableSynchronization.ON_READ,
+                                 aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA)
+
+    @tf.function
+    def set_as_moving_average_of(self, src_net):
+        beta, beta_nontrainable = 0.99, 0.0
+
+        for cw, sw in zip(self.weights, src_net.weights):
+            assert sw.shape == cw.shape
+            # print('{} <=> {}'.format(cw.name, sw.name))
+
+            if 'w_avg' in cw.name:
+                cw.assign(lerp(sw, cw, beta_nontrainable))
+            else:
+                cw.assign(lerp(sw, cw, beta))
+        return
+
+    def update_moving_average_of_w(self, w_broadcasted):
+        # compute average of current w
+        batch_avg = tf.reduce_mean(w_broadcasted[:, 0], axis=0)
+
+        # compute moving average of w and update(assign) w_avg
+        self.w_avg.assign(lerp(batch_avg, self.w_avg, self.w_ema_decay))
+        return
+
+    def style_mixing_regularization(self, latents1, labels, w_broadcasted1):
+        # get another w and broadcast it
+        latents2 = tf.random.normal(shape=tf.shape(latents1), dtype=tf.dtypes.float32)
+        dlatents2 = self.g_mapping([latents2, labels])
+        w_broadcasted2 = self.broadcast(dlatents2)
+
+        # find mixing limit index
+        # mixing_cutoff_index = tf.cond(
+        #     pred=tf.less(tf.random.uniform([], 0.0, 1.0), self.style_mixing_prob),
+        #     true_fn=lambda: tf.random.uniform([], 1, self.n_broadcast, dtype=tf.dtypes.int32),
+        #     false_fn=lambda: tf.constant(self.n_broadcast, dtype=tf.dtypes.int32))
+        if tf.random.uniform([], 0.0, 1.0) < self.style_mixing_prob:
+            mixing_cutoff_index = tf.random.uniform([], 1, self.n_broadcast, dtype=tf.dtypes.int32)
+        else:
+            mixing_cutoff_index = tf.constant(self.n_broadcast, dtype=tf.dtypes.int32)
+
+        # mix it
+        mixed_w_broadcasted = tf.where(
+            condition=tf.broadcast_to(self.mixing_layer_indices < mixing_cutoff_index, tf.shape(w_broadcasted1)),
+            x=w_broadcasted1,
+            y=w_broadcasted2)
+        return mixed_w_broadcasted
+
+    def truncation_trick(self, w_broadcasted, truncation_psi, truncation_cutoff=None):
+        ones = tf.ones_like(self.mixing_layer_indices, dtype=tf.float32)
+        tpsi = ones * truncation_psi
+        if truncation_cutoff is None:
+            truncation_coefs = tpsi
+        else:
+            indices = tf.range(self.n_broadcast)
+            truncation_coefs = tf.where(condition=tf.less(indices, truncation_cutoff), x=tpsi, y=ones)
+
+        truncated_w_broadcasted = lerp(self.w_avg, w_broadcasted, truncation_coefs)
+        return truncated_w_broadcasted
+
+    def call(self, inputs, ret_w_broadcasted=False, truncation_psi=1.0, truncation_cutoff=None, training=None, mask=None):
+        latents, labels = inputs
+
+        dlatents = self.g_mapping([latents, labels])
+        w_broadcasted = self.broadcast(dlatents)
+
+        if training:
+            self.update_moving_average_of_w(w_broadcasted)
+            w_broadcasted = self.style_mixing_regularization(latents, labels, w_broadcasted)
+
+        if not training:
+            w_broadcasted = self.truncation_trick(w_broadcasted, truncation_psi, truncation_cutoff)
+
+        image_out = self.synthesis(w_broadcasted)
+
+        if ret_w_broadcasted:
+            return image_out, w_broadcasted
+        else:
+            return image_out
+
+    def compute_output_shape(self, input_shape):
+        return input_shape[0][0], 3, self.resolutions[-1], self.resolutions[-1]

stylegan2_ref/image_proc.py

@@ -0,0 +1,120 @@
+import numpy as np
+import tensorflow as tf
+
+
+# ----------------------------------------------------------------------------
+# Primitive ops for manipulating 4D activation tensors.
+# The gradients of these are not necessary efficient or even meaningful.
+def _blur2d(x, f, normalize=True, flip=False, stride=1):
+    assert x.shape.ndims == 4 and all(dim is not None for dim in x.shape[1:])
+    assert isinstance(stride, int) and stride >= 1
+
+    # Finalize filter kernel.
+    f = np.array(f, dtype=np.float32)
+    if f.ndim == 1:
+        f = f[:, np.newaxis] * f[np.newaxis, :]
+    assert f.ndim == 2
+    if normalize:
+        f /= np.sum(f)
+    if flip:
+        f = f[::-1, ::-1]
+    f = f[:, :, np.newaxis, np.newaxis]
+    f = np.tile(f, [1, 1, int(x.shape[1]), 1])
+
+    # No-op => early exit.
+    if f.shape == (1, 1) and f[0, 0] == 1:
+        return x
+
+    # Convolve using depthwise_conv2d.
+    orig_dtype = x.dtype
+    x = tf.cast(x, tf.float32)  # tf.nn.depthwise_conv2d() doesn't support fp16
+    f = tf.constant(f, dtype=x.dtype, name='filter')
+    strides = [1, 1, stride, stride]
+    x = tf.nn.depthwise_conv2d(x, f, strides=strides, padding='SAME', data_format='NCHW')
+    x = tf.cast(x, orig_dtype)
+    return x
+
+
+def _upscale2d(x, factor=2, gain=1):
+    assert x.shape.ndims == 4 and all(dim is not None for dim in x.shape[1:])
+    assert isinstance(factor, int) and factor >= 1
+
+    # Apply gain.
+    if gain != 1:
+        x *= gain
+
+    # No-op => early exit.
+    if factor == 1:
+        return x
+
+    # Upscale using tf.tile().
+    s = x.shape
+    x = tf.reshape(x, [-1, s[1], s[2], 1, s[3], 1])
+    x = tf.tile(x, [1, 1, 1, factor, 1, factor])
+    x = tf.reshape(x, [-1, s[1], s[2] * factor, s[3] * factor])
+    return x
+
+
+def _downscale2d(x, factor=2, gain=1):
+    assert x.shape.ndims == 4 and all(dim is not None for dim in x.shape[1:])
+    assert isinstance(factor, int) and factor >= 1
+
+    # 2x2, float32 => downscale using _blur2d().
+    if factor == 2 and x.dtype == tf.float32:
+        f = [np.sqrt(gain) / factor] * factor
+        return _blur2d(x, f=f, normalize=False, stride=factor)
+
+    # Apply gain.
+    if gain != 1:
+        x *= gain
+
+    # No-op => early exit.
+    if factor == 1:
+        return x
+
+    # Large factor => downscale using tf.nn.avg_pool().
+    # NOTE: Requires tf_config['graph_options.place_pruned_graph']=True to work.
+    ksize = [1, 1, factor, factor]
+    return tf.nn.avg_pool(x, ksize=ksize, strides=ksize, padding='VALID', data_format='NCHW')
+
+
+# custom gradient functions
+def blur2d(x, f, normalize=True):
+    @tf.custom_gradient
+    def func(x):
+        y = _blur2d(x, f, normalize)
+
+        @tf.custom_gradient
+        def grad(dy):
+            dx = _blur2d(dy, f, normalize, flip=True)
+            return dx, lambda ddx: _blur2d(ddx, f, normalize)
+
+        return y, grad
+
+    return func(x)
+
+
+def upscale2d(x, factor=2):
+    @tf.custom_gradient
+    def func(x):
+        y = _upscale2d(x, factor)
+
+        @tf.custom_gradient
+        def grad(dy):
+            dx = _downscale2d(dy, factor, gain=factor ** 2)
+            return dx, lambda ddx: _upscale2d(ddx, factor)
+        return y, grad
+    return func(x)
+
+
+def downscale2d(x, factor=2):
+    @tf.custom_gradient
+    def func(x):
+        y = _downscale2d(x, factor)
+
+        @tf.custom_gradient
+        def grad(dy):
+            dx = _upscale2d(dy, factor, gain=1 / factor ** 2)
+            return dx, lambda ddx: _downscale2d(ddx, factor)
+        return y, grad
+    return func(x)

stylegan2_ref/upfirdn_2d.py

@@ -0,0 +1,118 @@
+import numpy as np
+import tensorflow as tf
+
+
+def setup_resample_kernel(k):
+    k = np.asarray(k, dtype=np.float32)
+    if k.ndim == 1:
+        k = np.outer(k, k)
+    k /= np.sum(k)
+    return k
+
+
+def upfirdn_ref(x, k, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1):
+    in_height, in_width = tf.shape(x)[1], tf.shape(x)[2]
+    minor_dim = tf.shape(x)[3]
+    kernel_h, kernel_w = k.shape
+
+    # Upsample (insert zeros).
+    x = tf.reshape(x, [-1, in_height, 1, in_width, 1, minor_dim])
+    x = tf.pad(x, [[0, 0], [0, 0], [0, up_y - 1], [0, 0], [0, up_x - 1], [0, 0]])
+    x = tf.reshape(x, [-1, in_height * up_y, in_width * up_x, minor_dim])
+
+    # Pad (crop if negative).
+    x = tf.pad(x, [
+        [0, 0], 
+        [tf.math.maximum(pad_y0, 0), tf.math.maximum(pad_y1, 0)], 
+        [tf.math.maximum(pad_x0, 0), tf.math.maximum(pad_x1, 0)], 
+        [0, 0]
+    ])
+    x = x[:, tf.math.maximum(-pad_y0, 0): tf.shape(x)[1] - tf.math.maximum(-pad_y1, 0),
+          tf.math.maximum(-pad_x0, 0): tf.shape(x)[2] - tf.math.maximum(-pad_x1, 0), :]
+
+    # Convolve with filter.
+    x = tf.transpose(x, [0, 3, 1, 2])
+    x = tf.reshape(x, [-1, 1, in_height * up_y + pad_y0 + pad_y1, in_width * up_x + pad_x0 + pad_x1])
+    w = tf.constant(k[::-1, ::-1, np.newaxis, np.newaxis], dtype=x.dtype)
+    x = tf.nn.conv2d(x, w, strides=[1, 1, 1, 1], padding='VALID', data_format='NCHW')
+    x = tf.reshape(x, [-1,
+                       minor_dim,
+                       in_height * up_y + pad_y0 + pad_y1 - kernel_h + 1,
+                       in_width * up_x + pad_x0 + pad_x1 - kernel_w + 1])
+    x = tf.transpose(x, [0, 2, 3, 1])
+
+    # Downsample (throw away pixels).
+    return x[:, ::down_y, ::down_x, :]
+
+
+def simple_upfirdn_2d(x, k, up=1, down=1, pad0=0, pad1=0):
+    output_channel = tf.shape(x)[1]
+    x = tf.reshape(x, [-1, tf.shape(x)[2], tf.shape(x)[3], 1])
+    x = upfirdn_ref(x, k,
+                    up_x=up, up_y=up, down_x=down, down_y=down, pad_x0=pad0, pad_x1=pad1, pad_y0=pad0, pad_y1=pad1)
+    x = tf.reshape(x, [-1, output_channel, tf.shape(x)[1], tf.shape(x)[2]])
+    return x
+
+
+def upsample_conv_2d(x, k, weight, factor, gain):
+    x_height, x_width = tf.shape(x)[2], tf.shape(x)[3]
+    w_height, w_width = tf.shape(weight)[0], tf.shape(weight)[1]
+    w_ic, w_oc = tf.shape(weight)[2], tf.shape(weight)[3]
+
+    # Setup filter kernel.
+    k = k * (gain * (factor ** 2))
+    p = (k.shape[0] - factor) - (w_width - 1)
+    pad0 = (p + 1) // 2 + factor - 1
+    pad1 = p // 2 + 1
+
+    # Determine data dimensions.
+    strides = [1, 1, factor, factor]
+    output_shape = [1, w_oc, (x_height - 1) * factor + w_height, (x_width - 1) * factor + w_width]
+    num_groups = tf.shape(x)[1] // w_ic
+
+    # Transpose weights.
+    weight = tf.reshape(weight, [w_height, w_width, w_ic, num_groups, -1])
+    weight = tf.transpose(weight[::-1, ::-1], [0, 1, 4, 3, 2])
+    weight = tf.reshape(weight, [w_height, w_width, -1, num_groups * w_ic])
+
+    # Execute.
+    x = tf.nn.conv2d_transpose(x, weight, output_shape, strides, padding='VALID', data_format='NCHW')
+    x = simple_upfirdn_2d(x, k, pad0=pad0, pad1=pad1)
+    return x
+
+
+def conv_downsample_2d(x, k, weight, factor, gain):
+    w_height, w_width = tf.shape(weight)[0], tf.shape(weight)[1]
+
+    # Setup filter kernel.
+    k = k * gain
+    p = (k.shape[0] - factor) + (w_width - 1)
+    pad0 = (p + 1) // 2
+    pad1 = p // 2
+
+    strides = [1, 1, factor, factor]
+    x = simple_upfirdn_2d(x, k, pad0=pad0, pad1=pad1)
+    x = tf.nn.conv2d(x, weight, strides, padding='VALID', data_format='NCHW')
+    return x
+
+
+def upsample_2d(x, k, factor, gain):
+    # Setup filter kernel.
+    k = k * (gain * (factor ** 2))
+    p = k.shape[0] - factor
+    pad0 = (p + 1) // 2 + factor - 1
+    pad1 = p // 2
+
+    x = simple_upfirdn_2d(x, k, up=factor, pad0=pad0, pad1=pad1)
+    return x
+
+
+def downsample_2d(x, k, factor, gain):
+    # Setup filter kernel.
+    k = k * gain
+    p = k.shape[0] - factor
+    pad0 = (p + 1) // 2
+    pad1 = p // 2
+
+    x = simple_upfirdn_2d(x, k, down=factor, pad0=pad0, pad1=pad1)
+    return x

stylegan2_ref/utils.py

@@ -0,0 +1,95 @@
+import numpy as np
+import tensorflow as tf
+
+
+def lerp(a, b, t):
+    out = a + (b - a) * t
+    return out
+
+
+def lerp_clip(a, b, t):
+    out = a + (b - a) * tf.clip_by_value(t, 0.0, 1.0)
+    return out
+
+
+def adjust_dynamic_range(images, range_in, range_out, out_dtype):
+    scale = (range_out[1] - range_out[0]) / (range_in[1] - range_in[0])
+    bias = range_out[0] - range_in[0] * scale
+    images = images * scale + bias
+    images = tf.clip_by_value(images, range_out[0], range_out[1])
+    images = tf.cast(images, dtype=out_dtype)
+    return images
+
+
+def random_flip_left_right_nchw(images):
+    s = tf.shape(images)
+    mask = tf.random.uniform([s[0], 1, 1, 1], 0.0, 1.0)
+    mask = tf.tile(mask, [1, s[1], s[2], s[3]])
+    images = tf.where(mask < 0.5, images, tf.reverse(images, axis=[3]))
+    return images
+
+
+def preprocess_fit_train_image(images, res):
+    images = adjust_dynamic_range(images, range_in=(0.0, 255.0), range_out=(-1.0, 1.0), out_dtype=tf.dtypes.float32)
+    images = random_flip_left_right_nchw(images)
+    images.set_shape([None, 3, res, res])
+    return images
+
+
+def postprocess_images(images):
+    images = adjust_dynamic_range(images, range_in=(-1.0, 1.0), range_out=(0.0, 255.0), out_dtype=tf.dtypes.float32)
+    images = tf.transpose(images, [0, 2, 3, 1])
+    images = tf.cast(images, dtype=tf.dtypes.uint8)
+    return images
+
+
+def merge_batch_images(images, res, rows, cols):
+    batch_size = images.shape[0]
+    assert rows * cols == batch_size
+    canvas = np.zeros(shape=[res * rows, res * cols, 3], dtype=np.uint8)
+    for row in range(rows):
+        y_start = row * res
+        for col in range(cols):
+            x_start = col * res
+            index = col + row * cols
+            canvas[y_start:y_start + res, x_start:x_start + res, :] = images[index, :, :, :]
+    return canvas
+
+
+def main():
+    # out_dtype = tf.dtypes.float32
+    # range_in = (0.0, 255.0)
+    # range_out = (-1.0, 1.0)
+    # # images = np.ones((2, 2), dtype=np.float32) * 255.0
+    # images1 = np.ones((2, 2), dtype=np.float32) * 0.0
+    # images2 = np.ones((2, 2), dtype=np.float32) * 127.5
+    # images3 = np.ones((2, 2), dtype=np.float32) * 255.0
+    # adjusted1 = adjust_dynamic_range(images1, range_in, range_out, out_dtype)
+    # adjusted2 = adjust_dynamic_range(images2, range_in, range_out, out_dtype)
+    # adjusted3 = adjust_dynamic_range(images3, range_in, range_out, out_dtype)
+    # print(adjusted1)
+    # print(adjusted2)
+    # print(adjusted3)
+    #
+    # out_dtype = tf.dtypes.uint8
+    # range_in = (-1.0, 1.0)
+    # range_out = (0.0, 255.0)
+    # images1 = np.ones((2, 2), dtype=np.float32) * -1.0
+    # images2 = np.ones((2, 2), dtype=np.float32) * 0.0
+    # images3 = np.ones((2, 2), dtype=np.float32) * 1.0
+    # adjusted1 = adjust_dynamic_range(images1, range_in, range_out, out_dtype)
+    # adjusted2 = adjust_dynamic_range(images2, range_in, range_out, out_dtype)
+    # adjusted3 = adjust_dynamic_range(images3, range_in, range_out, out_dtype)
+    # print(adjusted1)
+    # print(adjusted2)
+    # print(adjusted3)
+
+    batch_size = 8
+    res = 128
+    fake_images = np.ones(shape=(batch_size, res, res, 3), dtype=np.uint8)
+    merge_batch_images(fake_images, res, rows=4, cols=2)
+    return
+
+
+if __name__ == '__main__':
+    main()

tf_utils.py

@@ -0,0 +1,44 @@
+import tensorflow as tf
+
+from distutils.version import StrictVersion
+
+
+def allow_memory_growth():
+    gpus = tf.config.experimental.list_physical_devices('GPU')
+    if gpus:
+        try:
+            # Currently, memory growth needs to be the same across GPUs
+            for gpu in gpus:
+                tf.config.experimental.set_memory_growth(gpu, True)
+            logical_gpus = tf.config.experimental.list_logical_devices('GPU')
+            print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
+        except RuntimeError as e:
+            # Memory growth must be set before GPUs have been initialized
+            print(e)
+    return
+
+
+def split_gpu_for_testing(mem_in_gb=4):
+    gpus = tf.config.experimental.list_physical_devices('GPU')
+    if gpus:
+        try:
+            tf.config.experimental.set_virtual_device_configuration(
+                gpus[0],
+                [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024 * mem_in_gb),
+                 tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024 * mem_in_gb)]
+            )
+            logical_gpus = tf.config.experimental.list_logical_devices('GPU')
+            print(f'{len(gpus)} Physical GPUs, {len(logical_gpus)} Logical GPUs')
+        except RuntimeError as e:
+            print(e)
+    return
+
+
+def check_tf_version():
+    # check tensorflow version
+    tf_min_ver = '2.0.0'
+    cur_tf_ver = tf.__version__
+    print(f'Tensorflow version: {cur_tf_ver}')
+    if StrictVersion(cur_tf_ver) < StrictVersion(tf_min_ver):
+        raise ValueError(f'Need at least tf ver {tf_min_ver}')
+    return cur_tf_ver

train.py

@@ -0,0 +1,439 @@
+import os
+import time
+import argparse
+from cv2 import drawFrameAxes
+import numpy as np
+import tensorflow as tf
+
+from utils import str_to_bool
+from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
+from load_models import load_generator, load_discriminator
+from dataset_ffhq import get_ffhq_dataset
+from dataset import DatasetSimple
+from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg
+
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+#wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_2')
+
+
+def initiate_models(g_params, d_params, use_custom_cuda):
+    discriminator = load_discriminator(d_params, ckpt_dir=None, custom_cuda=use_custom_cuda)
+    generator = load_generator(g_params=g_params, is_g_clone=False, ckpt_dir=None, custom_cuda=use_custom_cuda)
+    g_clone = load_generator(g_params=g_params, is_g_clone=True, ckpt_dir=None, custom_cuda=use_custom_cuda)
+
+    # set initial g_clone weights same as generator
+    g_clone.set_weights(generator.get_weights())
+    return discriminator, generator, g_clone
+
+
+class Trainer(object):
+    def __init__(self, t_params, name):
+        self.cur_tf_ver = t_params['cur_tf_ver']
+        self.use_tf_function = t_params['use_tf_function']
+        self.use_custom_cuda = t_params['use_custom_cuda']
+        self.model_base_dir = t_params['model_base_dir']
+        self.global_batch_size = t_params['batch_size']
+        self.n_total_image = t_params['n_total_image']
+        self.max_steps = int(np.ceil(self.n_total_image / self.global_batch_size))
+        self.n_samples = min(t_params['batch_size'], t_params['n_samples'])
+        self.train_res = t_params['train_res']
+        self.print_step = 10
+        self.save_step = 1000
+        self.image_summary_step = 1000
+        self.reached_max_steps = False
+        self.log_template = '{:s}, {:s}, {:s}'.format(
+            'step {}: elapsed: {:.2f}s, d_loss: {:.3f}, g_loss: {:.3f}',
+            'd_gan_loss: {:.3f}, g_gan_loss: {:.3f}',
+            'r1_penalty: {:.3f}, pl_penalty: {:.3f}')
+
+        # copy network params
+        self.g_params = t_params['g_params']
+        self.d_params = t_params['d_params']
+
+        # set optimizer params
+        self.global_batch_scaler = 1.0 / float(self.global_batch_size)
+        self.r1_gamma = 10.0
+        self.g_opt = self.update_optimizer_params(t_params['g_opt'])
+        self.d_opt = self.update_optimizer_params(t_params['d_opt'])
+        self.pl_minibatch_shrink = 2
+        self.pl_weight = float(self.pl_minibatch_shrink)
+        self.pl_denorm = tf.math.rsqrt(float(self.train_res) * float(self.train_res))
+        self.pl_decay = 0.01
+        self.pl_mean = tf.Variable(initial_value=0.0, name='pl_mean', trainable=False,
+                                   synchronization=tf.VariableSynchronization.ON_READ,
+                                   aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA)
+
+        # create model: model and optimizer must be created under `strategy.scope`
+        self.discriminator, self.generator, self.g_clone = initiate_models(self.g_params,
+                                                                           self.d_params,
+                                                                           self.use_custom_cuda)
+
+        # set optimizers
+        self.d_optimizer = tf.keras.optimizers.Adam(self.d_opt['learning_rate'],
+                                                    beta_1=self.d_opt['beta1'],
+                                                    beta_2=self.d_opt['beta2'],
+                                                    epsilon=self.d_opt['epsilon'])
+        self.g_optimizer = tf.keras.optimizers.Adam(self.g_opt['learning_rate'],
+                                                    beta_1=self.g_opt['beta1'],
+                                                    beta_2=self.g_opt['beta2'],
+                                                    epsilon=self.g_opt['epsilon'])
+
+        # setup saving locations (object based savings)
+        self.ckpt_dir = os.path.join(self.model_base_dir, name)
+        self.ckpt = tf.train.Checkpoint(d_optimizer=self.d_optimizer,
+                                        g_optimizer=self.g_optimizer,
+                                        discriminator=self.discriminator,
+                                        generator=self.generator,
+                                        g_clone=self.g_clone,
+                                        pl_mean=self.pl_mean)
+        self.manager = tf.train.CheckpointManager(self.ckpt, self.ckpt_dir, max_to_keep=5)
+
+        # try to restore
+        self.ckpt.restore(self.manager.latest_checkpoint)
+        if self.manager.latest_checkpoint:
+            print('Restored from {}'.format(self.manager.latest_checkpoint))
+
+            # check if already trained in this resolution
+            restored_step = self.g_optimizer.iterations.numpy()
+            if restored_step >= self.max_steps:
+                print('Already reached max steps {}/{}'.format(restored_step, self.max_steps))
+                self.reached_max_steps = True
+                return
+        else:
+            print('Not restoring from saved checkpoint')
+
+    @staticmethod
+    def update_optimizer_params(params):
+        params_copy = params.copy()
+        mb_ratio = params_copy['reg_interval'] / (params_copy['reg_interval'] + 1)
+        params_copy['learning_rate'] = params_copy['learning_rate'] * mb_ratio
+        params_copy['beta1'] = params_copy['beta1'] ** mb_ratio
+        params_copy['beta2'] = params_copy['beta2'] ** mb_ratio
+        return params_copy
+
+    def d_train_step(self, dist_inputs):
+        real_images, condition = dist_inputs[0][0], dist_inputs[0][1]
+
+        with tf.GradientTape() as d_tape:
+            # compute losses
+            d_loss, cond_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'], labels=condition)
+
+            # scale loss
+            d_loss = tf.reduce_sum(d_loss) * self.global_batch_scaler
+            cond_loss = tf.reduce_sum(cond_loss) * self.global_batch_scaler
+
+            d_loss = d_loss + cond_loss
+
+        d_gradients = d_tape.gradient(d_loss, self.discriminator.trainable_variables)
+        self.d_optimizer.apply_gradients(zip(d_gradients, self.discriminator.trainable_variables))
+        return d_loss, cond_loss
+
+    def d_train_step_reg(self, dist_inputs):
+        real_images, condition = dist_inputs[0][0], dist_inputs[0][1]
+
+        with tf.GradientTape() as d_tape:
+            # compute losses
+            d_gan_loss, cond_loss, r1_penalty = d_logistic_r1_reg(real_images, self.generator, self.discriminator,
+                                                       self.g_params['z_dim'], labels=condition)
+            r1_penalty = r1_penalty * (0.5 * self.r1_gamma) * self.d_opt['reg_interval']
+
+            # scale losses
+            r1_penalty = tf.reduce_sum(r1_penalty) * self.global_batch_scaler
+            d_gan_loss = tf.reduce_sum(d_gan_loss) * self.global_batch_scaler
+            cond_loss = tf.reduce_sum(cond_loss) * self.global_batch_scaler
+
+            # combine
+            d_loss = d_gan_loss + r1_penalty + cond_loss
+
+        d_gradients = d_tape.gradient(d_loss, self.discriminator.trainable_variables)
+        self.d_optimizer.apply_gradients(zip(d_gradients, self.discriminator.trainable_variables))
+        return d_loss, d_gan_loss, r1_penalty, cond_loss
+
+    def g_train_step(self, dist_inputs):
+        real_images, condition = dist_inputs[0][0], dist_inputs[0][1]
+
+        with tf.GradientTape() as g_tape:
+            # compute losses
+            g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], labels=condition)
+
+            # scale loss
+            g_loss = tf.reduce_sum(g_loss) * self.global_batch_scaler
+
+        g_gradients = g_tape.gradient(g_loss, self.generator.trainable_variables)
+        self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_variables))
+        return g_loss
+
+    def g_train_step_reg(self, dist_inputs):
+        real_images, condition = dist_inputs[0][0], dist_inputs[0][1]
+
+        with tf.GradientTape() as g_tape:
+            # compute losses
+            g_gan_loss, pl_penalty = g_logistic_ns_pathreg(real_images, self.generator, self.discriminator,
+                                                           self.g_params['z_dim'], self.pl_mean,
+                                                           self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, labels=condition)
+            pl_penalty = pl_penalty * self.pl_weight * self.g_opt['reg_interval']
+
+            # scale loss
+            pl_penalty = tf.reduce_sum(pl_penalty) * self.global_batch_scaler
+            g_gan_loss = tf.reduce_sum(g_gan_loss) * self.global_batch_scaler
+
+            # combine
+            g_loss = g_gan_loss + pl_penalty
+
+        g_gradients = g_tape.gradient(g_loss, self.generator.trainable_variables)
+        self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_variables))
+        return g_loss, g_gan_loss, pl_penalty
+
+    def train(self, dist_datasets, strategy):
+        # @tf.function
+        def dist_d_train_step(inputs):
+            per_replica_losses = strategy.run(fn=self.d_train_step, args=(inputs,))
+            mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
+            mean_cond_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
+            return mean_d_loss, mean_cond_loss
+
+        # @tf.function
+        def dist_d_train_step_reg(inputs):
+            per_replica_losses = strategy.run(fn=self.d_train_step_reg, args=(inputs,))
+            mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
+            mean_d_gan_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
+            mean_r1_penalty = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
+            mean_cond_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None)
+            return mean_d_loss, mean_d_gan_loss, mean_r1_penalty, mean_cond_loss
+
+        # @tf.function
+        def dist_g_train_step(inputs):
+            per_replica_losses = strategy.run(fn=self.g_train_step, args=(inputs,))
+            mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
+            return mean_g_loss
+
+        # @tf.function
+        def dist_g_train_step_reg(inputs):
+            per_replica_losses = strategy.run(fn=self.g_train_step_reg, args=(inputs,))
+            mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
+            mean_g_gan_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
+            mean_pl_penalty = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
+            return mean_g_loss, mean_g_gan_loss, mean_pl_penalty
+
+        @tf.function
+        def dist_gen_samples(dist_inputs):
+            per_replica_samples = strategy.run(self.gen_samples, args=(dist_inputs,))
+            return per_replica_samples
+
+        if self.reached_max_steps:
+            return
+
+        # start actual training
+        print('Start Training')
+
+        # setup tensorboards
+        train_summary_writer = tf.summary.create_file_writer(self.ckpt_dir)
+
+        # loss metrics
+        metric_d_loss = tf.keras.metrics.Mean('d_loss', dtype=tf.float32)
+        metric_g_loss = tf.keras.metrics.Mean('g_loss', dtype=tf.float32)
+        metric_d_gan_loss = tf.keras.metrics.Mean('d_gan_loss', dtype=tf.float32)
+        metric_g_gan_loss = tf.keras.metrics.Mean('g_gan_loss', dtype=tf.float32)
+        metric_r1_penalty = tf.keras.metrics.Mean('r1_penalty', dtype=tf.float32)
+        metric_pl_penalty = tf.keras.metrics.Mean('pl_penalty', dtype=tf.float32)
+        metric_cond_loss = tf.keras.metrics.Mean('cond_loss', dtype=tf.float32)
+
+        # start training
+        zero = tf.constant(0.0, dtype=tf.float32)
+        print('max_steps: {}'.format(self.max_steps))
+        t_start = time.time()
+        step = self.g_optimizer.iterations.numpy()
+        for real_images in dist_datasets:
+
+            # d train step
+            if (step + 1) % self.d_opt['reg_interval'] == 0:
+                d_loss, d_gan_loss, r1_penalty, cond_loss = dist_d_train_step_reg((real_images, ))
+            else:
+                d_loss, cond_loss = dist_d_train_step((real_images, ))
+                d_gan_loss = d_loss
+                r1_penalty = zero
+
+            # g train step
+            if (step + 1) % self.g_opt['reg_interval'] == 0:
+                g_loss, g_gan_loss, pl_penalty = dist_g_train_step_reg((real_images, ))
+            else:
+                g_loss = dist_g_train_step((real_images,))
+                g_gan_loss = g_loss
+                pl_penalty = zero
+
+            # update g_clone
+            self.g_clone.set_as_moving_average_of(self.generator)
+
+            # update metrics
+            metric_d_loss(d_loss)
+            metric_g_loss(g_loss)
+            metric_d_gan_loss(d_gan_loss)
+            metric_g_gan_loss(g_gan_loss)
+            metric_r1_penalty(r1_penalty)
+            metric_pl_penalty(pl_penalty)
+            metric_cond_loss(cond_loss)
+
+            # get current step
+            step += 1
+
+            # save to tensorboard
+            with train_summary_writer.as_default():
+                tf.summary.scalar('d_loss', metric_d_loss.result(), step=step)
+                tf.summary.scalar('g_loss', metric_g_loss.result(), step=step)
+                tf.summary.scalar('d_gan_loss', metric_d_gan_loss.result(), step=step)
+                tf.summary.scalar('g_gan_loss', metric_g_gan_loss.result(), step=step)
+                tf.summary.scalar('r1_penalty', metric_r1_penalty.result(), step=step)
+                tf.summary.scalar('pl_penalty', metric_pl_penalty.result(), step=step)
+                tf.summary.scalar('cond_loss', metric_cond_loss.result(), step=step)
+
+            # save every self.save_step
+            if step % self.save_step == 0:
+                self.manager.save(checkpoint_number=step)
+
+            # save every self.image_summary_step
+            if step % self.image_summary_step == 0:
+                # add summary image
+                test_z = tf.random.normal(shape=(self.n_samples, self.g_params['z_dim']), dtype=tf.dtypes.float32)
+                test_labels = tf.ones((self.n_samples, self.g_params['labels_dim']), dtype=tf.dtypes.float32)
+                summary_image = dist_gen_samples((test_z, test_labels))
+
+                # convert to tensor image
+                summary_image = self.convert_per_replica_image(summary_image, strategy)
+
+                with train_summary_writer.as_default():
+                    tf.summary.image('images', summary_image, step=step)
+
+                # train_summary_writer.flush()
+
+            # print every self.print_steps
+            if step % self.print_step == 0:
+                elapsed = time.time() - t_start
+                print(self.log_template.format(step, elapsed, d_loss.numpy(), g_loss.numpy(),
+                                               d_gan_loss.numpy(), g_gan_loss.numpy(),
+                                               r1_penalty.numpy(), pl_penalty.numpy()))
+
+                # reset timer
+                t_start = time.time()
+
+            # check exit status
+            if step >= self.max_steps:
+                break
+
+        # save last checkpoint
+        step = self.g_optimizer.iterations.numpy()
+        self.manager.save(checkpoint_number=step)
+        return
+
+    def gen_samples(self, inputs):
+        test_z, test_labels = inputs
+
+        # run networks
+        fake_images_05 = self.g_clone([test_z, test_labels], truncation_psi=0.5, training=False)
+        fake_images_07 = self.g_clone([test_z, test_labels], truncation_psi=0.7, training=False)
+
+        # merge on batch dimension: [n_samples, 3, out_res, 2 * out_res]
+        final_image = tf.concat([fake_images_05, fake_images_07], axis=2)
+        return final_image
+
+    @staticmethod
+    def convert_per_replica_image(nchw_per_replica_images, strategy):
+        as_tensor = tf.concat(strategy.experimental_local_results(nchw_per_replica_images), axis=0)
+        as_tensor = tf.transpose(as_tensor, perm=[0, 2, 3, 1])
+        as_tensor = (tf.clip_by_value(as_tensor, -1.0, 1.0) + 1.0) * 127.5
+        as_tensor = tf.cast(as_tensor, tf.uint8)
+        return as_tensor
+
+
+def filter_resolutions_featuremaps(resolutions, featuremaps, res):
+    index = resolutions.index(res)
+    filtered_resolutions = resolutions[:index + 1]
+    filtered_featuremaps = featuremaps[:index + 1]
+    return filtered_resolutions, filtered_featuremaps
+
+
+def main():
+    # global program arguments parser
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--allow_memory_growth', type=str_to_bool, nargs='?', const=True, default=True)
+    parser.add_argument('--debug_split_gpu', type=str_to_bool, nargs='?', const=True, default=False)
+    parser.add_argument('--use_tf_function', type=str_to_bool, nargs='?', const=True, default=True)
+    parser.add_argument('--use_custom_cuda', type=str_to_bool, nargs='?', const=True, default=True)
+    parser.add_argument('--model_base_dir', default='./models', type=str)
+    parser.add_argument('--tfrecord_dir', default='./tfrecords', type=str)
+    parser.add_argument('--train_res', default=128, type=int)
+    parser.add_argument('--shuffle_buffer_size', default=1000, type=int)
+    parser.add_argument('--batch_size_per_replica', default=1, type=int)
+    args = vars(parser.parse_args())
+
+    # check tensorflow version
+    cur_tf_ver = check_tf_version()
+
+    # GPU environment settings
+    if args['allow_memory_growth']:
+        allow_memory_growth()
+    if args['debug_split_gpu']:
+        split_gpu_for_testing(mem_in_gb=4.5)
+
+    # network params
+    resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
+    # featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
+    featuremaps = [64, 64, 64, 64, 64, 32, 16, 8, 4]
+    train_resolutions, train_featuremaps = filter_resolutions_featuremaps(resolutions, featuremaps, args['train_res'])
+    g_params = {
+        'z_dim': 16,
+        'w_dim': 16,
+        'labels_dim': 83,
+        'n_mapping': 8,
+        'resolutions': train_resolutions,
+        'featuremaps': train_featuremaps,
+    }
+    d_params = {
+        'labels_dim': 83,
+        'resolutions': train_resolutions,
+        'featuremaps': train_featuremaps,
+    }
+
+    # prepare distribute strategy
+    strategy = tf.distribute.MirroredStrategy()
+    global_batch_size = args['batch_size_per_replica'] * strategy.num_replicas_in_sync
+
+    # prepare dataset
+    # dataset = get_ffhq_dataset(args['tfrecord_dir'], args['train_res'], batch_size=global_batch_size, epochs=None)
+    dataset = DatasetSimple('/workspace/Dataset/stylegan').batch(global_batch_size)#.shuffle(buffer_size=1000).prefetch(buffer_size=tf.data.experimental.AUTOTUNE).cache()
+    options = tf.data.Options()
+    options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
+    dataset = dataset.with_options(options)
+
+    with strategy.scope():
+        # distribute dataset        
+        dist_dataset = strategy.experimental_distribute_dataset(dataset)
+
+        # training parameters
+        training_parameters = {
+            # global params
+            'cur_tf_ver': cur_tf_ver,
+            'use_tf_function': args['use_tf_function'],
+            'use_custom_cuda': args['use_custom_cuda'],
+            'model_base_dir': args['model_base_dir'],
+
+            # network params
+            'g_params': g_params,
+            'd_params': d_params,
+
+            # training params
+            'g_opt': {'learning_rate': 0.002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 8},
+            'd_opt': {'learning_rate': 0.002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 16},
+            'batch_size': global_batch_size,
+            'n_total_image': 25000000,
+            'n_samples': 12,
+            'train_res': args['train_res'],
+        }
+
+        trainer = Trainer(training_parameters, name=f'stylegan2-test-4-{args["train_res"]}x{args["train_res"]}')
+        trainer.train(dist_dataset, strategy)
+    return
+
+
+if __name__ == '__main__':
+    main()

train_memtest.py

@@ -0,0 +1,480 @@
+import os
+import time
+import argparse
+import linecache
+import tracemalloc
+from cv2 import drawFrameAxes
+import numpy as np
+import tensorflow as tf
+
+from utils import str_to_bool
+from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
+from load_models import load_generator, load_discriminator
+from dataset_ffhq import get_ffhq_dataset
+from dataset import DatasetSimple
+from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg
+
+def display_top(snapshot, key_type='lineno', limit=100):
+    snapshot = snapshot.filter_traces((
+        tracemalloc.Filter(False, "<frozen importlib._bootstrap>"),
+        tracemalloc.Filter(False, "<unknown>"),
+    ))
+    top_stats = snapshot.statistics(key_type)
+
+    print("Top %s lines" % limit)
+    for index, stat in enumerate(top_stats[:limit], 1):
+        frame = stat.traceback[0]
+        print("#%s: %s:%s: %.1f KiB"
+              % (index, frame.filename, frame.lineno, stat.size / 1024))
+        line = linecache.getline(frame.filename, frame.lineno).strip()
+        if line:
+            print('    %s' % line)
+
+    other = top_stats[limit:]
+    if other:
+        size = sum(stat.size for stat in other)
+        print("%s other: %.1f KiB" % (len(other), size / 1024))
+    total = sum(stat.size for stat in top_stats)
+    print("Total allocated size: %.1f KiB" % (total / 1024))
+
+    
+tracemalloc.start()
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+
+
+def initiate_models(g_params, d_params, use_custom_cuda):
+    discriminator = load_discriminator(d_params, ckpt_dir=None, custom_cuda=use_custom_cuda)
+    generator = load_generator(g_params=g_params, is_g_clone=False, ckpt_dir=None, custom_cuda=use_custom_cuda)
+    g_clone = load_generator(g_params=g_params, is_g_clone=True, ckpt_dir=None, custom_cuda=use_custom_cuda)
+
+    # set initial g_clone weights same as generator
+    g_clone.set_weights(generator.get_weights())
+    return discriminator, generator, g_clone
+
+
+class Trainer(object):
+    def __init__(self, t_params, name):
+        self.cur_tf_ver = t_params['cur_tf_ver']
+        self.use_tf_function = t_params['use_tf_function']
+        self.use_custom_cuda = t_params['use_custom_cuda']
+        self.model_base_dir = t_params['model_base_dir']
+        self.global_batch_size = t_params['batch_size']
+        self.n_total_image = t_params['n_total_image']
+        self.max_steps = int(np.ceil(self.n_total_image / self.global_batch_size))
+        self.n_samples = min(t_params['batch_size'], t_params['n_samples'])
+        self.train_res = t_params['train_res']
+        self.print_step = 10
+        self.save_step = 1000
+        self.image_summary_step = 10000
+        self.reached_max_steps = False
+        self.log_template = '{:s}, {:s}, {:s}'.format(
+            'step {}: elapsed: {:.2f}s, d_loss: {:.3f}, g_loss: {:.3f}',
+            'd_gan_loss: {:.3f}, g_gan_loss: {:.3f}',
+            'r1_penalty: {:.3f}, pl_penalty: {:.3f}')
+
+        # copy network params
+        self.g_params = t_params['g_params']
+        self.d_params = t_params['d_params']
+
+        # set optimizer params
+        self.global_batch_scaler = 1.0 / float(self.global_batch_size)
+        self.r1_gamma = 10.0
+        self.g_opt = self.update_optimizer_params(t_params['g_opt'])
+        self.d_opt = self.update_optimizer_params(t_params['d_opt'])
+        self.pl_minibatch_shrink = 2
+        self.pl_weight = float(self.pl_minibatch_shrink)
+        self.pl_denorm = tf.math.rsqrt(float(self.train_res) * float(self.train_res))
+        self.pl_decay = 0.01
+        self.pl_mean = tf.Variable(initial_value=0.0, name='pl_mean', trainable=False,
+                                   synchronization=tf.VariableSynchronization.ON_READ,
+                                   aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA)
+
+        # create model: model and optimizer must be created under `strategy.scope`
+        self.discriminator, self.generator, self.g_clone = initiate_models(self.g_params,
+                                                                           self.d_params,
+                                                                           self.use_custom_cuda)
+
+        # set optimizers
+        self.d_optimizer = tf.keras.optimizers.Adam(self.d_opt['learning_rate'],
+                                                    beta_1=self.d_opt['beta1'],
+                                                    beta_2=self.d_opt['beta2'],
+                                                    epsilon=self.d_opt['epsilon'])
+        self.g_optimizer = tf.keras.optimizers.Adam(self.g_opt['learning_rate'],
+                                                    beta_1=self.g_opt['beta1'],
+                                                    beta_2=self.g_opt['beta2'],
+                                                    epsilon=self.g_opt['epsilon'])
+
+        # setup saving locations (object based savings)
+        self.ckpt_dir = os.path.join(self.model_base_dir, name)
+        self.ckpt = tf.train.Checkpoint(d_optimizer=self.d_optimizer,
+                                        g_optimizer=self.g_optimizer,
+                                        discriminator=self.discriminator,
+                                        generator=self.generator,
+                                        g_clone=self.g_clone,
+                                        pl_mean=self.pl_mean)
+        self.manager = tf.train.CheckpointManager(self.ckpt, self.ckpt_dir, max_to_keep=1)
+
+        # try to restore
+        self.ckpt.restore(self.manager.latest_checkpoint)
+        if self.manager.latest_checkpoint:
+            print('Restored from {}'.format(self.manager.latest_checkpoint))
+
+            # check if already trained in this resolution
+            restored_step = self.g_optimizer.iterations.numpy()
+            if restored_step >= self.max_steps:
+                print('Already reached max steps {}/{}'.format(restored_step, self.max_steps))
+                self.reached_max_steps = True
+                return
+        else:
+            print('Not restoring from saved checkpoint')
+
+    @staticmethod
+    def update_optimizer_params(params):
+        params_copy = params.copy()
+        mb_ratio = params_copy['reg_interval'] / (params_copy['reg_interval'] + 1)
+        params_copy['learning_rate'] = params_copy['learning_rate'] * mb_ratio
+        params_copy['beta1'] = params_copy['beta1'] ** mb_ratio
+        params_copy['beta2'] = params_copy['beta2'] ** mb_ratio
+        return params_copy
+
+    def d_train_step(self, dist_inputs):
+        real_images = dist_inputs[0]
+
+        with tf.GradientTape() as d_tape:
+            # compute losses
+            d_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'])
+
+            # scale loss
+            d_loss = tf.reduce_sum(d_loss) * self.global_batch_scaler
+
+        d_gradients = d_tape.gradient(d_loss, self.discriminator.trainable_variables)
+        self.d_optimizer.apply_gradients(zip(d_gradients, self.discriminator.trainable_variables))
+        return d_loss
+
+    def d_train_step_reg(self, dist_inputs):
+        real_images = dist_inputs[0]
+
+        with tf.GradientTape() as d_tape:
+            # compute losses
+            d_gan_loss, r1_penalty = d_logistic_r1_reg(real_images, self.generator, self.discriminator,
+                                                       self.g_params['z_dim'])
+            r1_penalty = r1_penalty * (0.5 * self.r1_gamma) * self.d_opt['reg_interval']
+
+            # scale losses
+            r1_penalty = tf.reduce_sum(r1_penalty) * self.global_batch_scaler
+            d_gan_loss = tf.reduce_sum(d_gan_loss) * self.global_batch_scaler
+
+            # combine
+            d_loss = d_gan_loss + r1_penalty
+
+        d_gradients = d_tape.gradient(d_loss, self.discriminator.trainable_variables)
+        self.d_optimizer.apply_gradients(zip(d_gradients, self.discriminator.trainable_variables))
+        return d_loss, d_gan_loss, r1_penalty
+
+    def g_train_step(self, dist_inputs):
+        real_images = dist_inputs[0]
+
+        with tf.GradientTape() as g_tape:
+            # compute losses
+            g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'])
+
+            # scale loss
+            g_loss = tf.reduce_sum(g_loss) * self.global_batch_scaler
+
+        g_gradients = g_tape.gradient(g_loss, self.generator.trainable_variables)
+        self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_variables))
+        return g_loss
+
+    def g_train_step_reg(self, dist_inputs):
+        real_images = dist_inputs[0]
+
+        with tf.GradientTape() as g_tape:
+            # compute losses
+            g_gan_loss, pl_penalty = g_logistic_ns_pathreg(real_images, self.generator, self.discriminator,
+                                                           self.g_params['z_dim'], self.pl_mean,
+                                                           self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay)
+            pl_penalty = pl_penalty * self.pl_weight * self.g_opt['reg_interval']
+
+            # scale loss
+            pl_penalty = tf.reduce_sum(pl_penalty) * self.global_batch_scaler
+            g_gan_loss = tf.reduce_sum(g_gan_loss) * self.global_batch_scaler
+
+            # combine
+            g_loss = g_gan_loss + pl_penalty
+
+        g_gradients = g_tape.gradient(g_loss, self.generator.trainable_variables)
+        self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_variables))
+        return g_loss, g_gan_loss, pl_penalty
+
+#    @tf.function
+    def train(self, dist_datasets, strategy):
+        def dist_d_train_step(inputs):
+            if self.cur_tf_ver == '2.0.0':
+                per_replica_losses = strategy.experimental_run_v2(fn=self.d_train_step, args=(inputs,))
+            else:
+                per_replica_losses = strategy.run(fn=self.d_train_step, args=(inputs,))
+            mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
+            return mean_d_loss
+
+        def dist_d_train_step_reg(inputs):
+            if self.cur_tf_ver == '2.0.0':
+                per_replica_losses = strategy.experimental_run_v2(fn=self.d_train_step_reg, args=(inputs,))
+            else:
+                per_replica_losses = strategy.run(fn=self.d_train_step_reg, args=(inputs,))
+            mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
+            mean_d_gan_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
+            mean_r1_penalty = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
+            return mean_d_loss, mean_d_gan_loss, mean_r1_penalty
+
+        def dist_g_train_step(inputs):
+            if self.cur_tf_ver == '2.0.0':
+                per_replica_losses = strategy.experimental_run_v2(fn=self.g_train_step, args=(inputs,))
+            else:
+                per_replica_losses = strategy.run(fn=self.g_train_step, args=(inputs,))
+            mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
+            return mean_g_loss
+
+        def dist_g_train_step_reg(inputs):
+            if self.cur_tf_ver == '2.0.0':
+                per_replica_losses = strategy.experimental_run_v2(fn=self.g_train_step_reg, args=(inputs,))
+            else:
+                per_replica_losses = strategy.run(fn=self.g_train_step_reg, args=(inputs,))
+            mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
+            mean_g_gan_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
+            mean_pl_penalty = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
+            return mean_g_loss, mean_g_gan_loss, mean_pl_penalty
+
+        def dist_gen_samples(dist_inputs):
+            if self.cur_tf_ver == '2.0.0':
+                per_replica_samples = strategy.experimental_run_v2(self.gen_samples, args=(dist_inputs,))
+            else:
+                per_replica_samples = strategy.run(self.gen_samples, args=(dist_inputs,))
+            return per_replica_samples
+
+        # wrap with tf.function
+        if self.use_tf_function:
+            dist_d_train_step = tf.function(dist_d_train_step)
+            dist_g_train_step = tf.function(dist_g_train_step)
+            dist_d_train_step_reg = tf.function(dist_d_train_step_reg)
+            dist_g_train_step_reg = tf.function(dist_g_train_step_reg)
+            dist_gen_samples = tf.function(dist_gen_samples)
+
+        if self.reached_max_steps:
+            return
+
+        # start actual training
+        print('Start Training')
+
+        # setup tensorboards
+        train_summary_writer = tf.summary.create_file_writer(self.ckpt_dir)
+
+        # loss metrics
+        metric_d_loss = tf.keras.metrics.Mean('d_loss', dtype=tf.float32)
+        metric_g_loss = tf.keras.metrics.Mean('g_loss', dtype=tf.float32)
+        metric_d_gan_loss = tf.keras.metrics.Mean('d_gan_loss', dtype=tf.float32)
+        metric_g_gan_loss = tf.keras.metrics.Mean('g_gan_loss', dtype=tf.float32)
+        metric_r1_penalty = tf.keras.metrics.Mean('r1_penalty', dtype=tf.float32)
+        metric_pl_penalty = tf.keras.metrics.Mean('pl_penalty', dtype=tf.float32)
+
+        # start training
+        zero = tf.constant(0.0, dtype=tf.float32)
+        print('max_steps: {}'.format(self.max_steps))
+        t_start = time.time()
+        for real_images in dist_datasets:
+            step = self.g_optimizer.iterations.numpy()
+
+            # d train step
+            if (step + 1) % self.d_opt['reg_interval'] == 0:
+                d_loss, d_gan_loss, r1_penalty = dist_d_train_step_reg((real_images, ))
+            else:
+                d_loss = dist_d_train_step((real_images, ))
+                d_gan_loss = d_loss
+                r1_penalty = zero
+
+            # g train step
+            if (step + 1) % self.g_opt['reg_interval'] == 0:
+                g_loss, g_gan_loss, pl_penalty = dist_g_train_step_reg((real_images, ))
+            else:
+                g_loss = dist_g_train_step((real_images,))
+                g_gan_loss = g_loss
+                pl_penalty = zero
+
+            # update g_clone
+            self.g_clone.set_as_moving_average_of(self.generator)
+
+            # update metrics
+            metric_d_loss(d_loss)
+            metric_g_loss(g_loss)
+            metric_d_gan_loss(d_gan_loss)
+            metric_g_gan_loss(g_gan_loss)
+            metric_r1_penalty(r1_penalty)
+            metric_pl_penalty(pl_penalty)
+
+            # get current step
+            step = self.g_optimizer.iterations.numpy()
+
+            # save to tensorboard
+            with train_summary_writer.as_default():
+                tf.summary.scalar('d_loss', metric_d_loss.result(), step=step)
+                tf.summary.scalar('g_loss', metric_g_loss.result(), step=step)
+                tf.summary.scalar('d_gan_loss', metric_d_gan_loss.result(), step=step)
+                tf.summary.scalar('g_gan_loss', metric_g_gan_loss.result(), step=step)
+                tf.summary.scalar('r1_penalty', metric_r1_penalty.result(), step=step)
+                tf.summary.scalar('pl_penalty', metric_pl_penalty.result(), step=step)
+
+            # save every self.save_step
+            if step % self.save_step == 0:
+                self.manager.save(checkpoint_number=step)
+
+            # save every self.image_summary_step
+            if step % self.image_summary_step == 0:
+                # add summary image
+                test_z = tf.random.normal(shape=(self.n_samples, self.g_params['z_dim']), dtype=tf.dtypes.float32)
+                test_labels = tf.ones((self.n_samples, self.g_params['labels_dim']), dtype=tf.dtypes.float32)
+                summary_image = dist_gen_samples((test_z, test_labels))
+
+                # convert to tensor image
+                summary_image = self.convert_per_replica_image(summary_image, strategy)
+
+                with train_summary_writer.as_default():
+                    tf.summary.image('images', summary_image, step=step)
+
+            # print every self.print_steps
+            if step % self.print_step == 0:
+                elapsed = time.time() - t_start
+                print(self.log_template.format(step, elapsed, d_loss.numpy(), g_loss.numpy(),
+                                               d_gan_loss.numpy(), g_gan_loss.numpy(),
+                                               r1_penalty.numpy(), pl_penalty.numpy()))
+
+                # reset timer
+                t_start = time.time()
+
+            # check exit status
+            if step >= self.max_steps:
+                break
+
+        # save last checkpoint
+        step = self.g_optimizer.iterations.numpy()
+        self.manager.save(checkpoint_number=step)
+        return
+
+    def gen_samples(self, inputs):
+        test_z, test_labels = inputs
+
+        # run networks
+        fake_images_05 = self.g_clone([test_z, test_labels], truncation_psi=0.5, training=False)
+        fake_images_07 = self.g_clone([test_z, test_labels], truncation_psi=0.7, training=False)
+
+        # merge on batch dimension: [n_samples, 3, out_res, 2 * out_res]
+        final_image = tf.concat([fake_images_05, fake_images_07], axis=2)
+        return final_image
+
+    @staticmethod
+    def convert_per_replica_image(nchw_per_replica_images, strategy):
+        as_tensor = tf.concat(strategy.experimental_local_results(nchw_per_replica_images), axis=0)
+        as_tensor = tf.transpose(as_tensor, perm=[0, 2, 3, 1])
+        as_tensor = (tf.clip_by_value(as_tensor, -1.0, 1.0) + 1.0) * 127.5
+        as_tensor = tf.cast(as_tensor, tf.uint8)
+        return as_tensor
+
+
+def filter_resolutions_featuremaps(resolutions, featuremaps, res):
+    index = resolutions.index(res)
+    filtered_resolutions = resolutions[:index + 1]
+    filtered_featuremaps = featuremaps[:index + 1]
+    return filtered_resolutions, filtered_featuremaps
+
+
+def main():
+    # global program arguments parser
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--allow_memory_growth', type=str_to_bool, nargs='?', const=True, default=True)
+    parser.add_argument('--debug_split_gpu', type=str_to_bool, nargs='?', const=True, default=False)
+    parser.add_argument('--use_tf_function', type=str_to_bool, nargs='?', const=True, default=True)
+    parser.add_argument('--use_custom_cuda', type=str_to_bool, nargs='?', const=True, default=False)
+    parser.add_argument('--model_base_dir', default='./models', type=str)
+    parser.add_argument('--tfrecord_dir', default='./tfrecords', type=str)
+    parser.add_argument('--train_res', default=512, type=int)
+    parser.add_argument('--shuffle_buffer_size', default=1000, type=int)
+    parser.add_argument('--batch_size_per_replica', default=10, type=int)
+    args = vars(parser.parse_args())
+
+    # check tensorflow version
+    cur_tf_ver = check_tf_version()
+
+    # GPU environment settings
+    if args['allow_memory_growth']:
+        allow_memory_growth()
+    if args['debug_split_gpu']:
+        split_gpu_for_testing(mem_in_gb=4.5)
+
+    # network params
+    resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
+    featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
+    train_resolutions, train_featuremaps = filter_resolutions_featuremaps(resolutions, featuremaps, args['train_res'])
+    g_params = {
+        'z_dim': 512,
+        'w_dim': 512,
+        'labels_dim': 0,
+        'n_mapping': 8,
+        'resolutions': train_resolutions,
+        'featuremaps': train_featuremaps,
+    }
+    d_params = {
+        'labels_dim': 0,
+        'resolutions': train_resolutions,
+        'featuremaps': train_featuremaps,
+    }
+
+    # prepare distribute strategy
+    strategy = tf.distribute.MirroredStrategy()
+    global_batch_size = args['batch_size_per_replica'] * strategy.num_replicas_in_sync
+
+    # prepare dataset
+    # dataset = get_ffhq_dataset(args['tfrecord_dir'], args['train_res'], batch_size=global_batch_size, epochs=None)
+    dataset = DatasetSimple('/home/workspace/local-data/stylegan/npz').shuffle(buffer_size=100).repeat(None).batch(global_batch_size).prefetch(buffer_size=100).cache()
+
+    with strategy.scope():
+        # distribute dataset
+        #dist_dataset = strategy.experimental_distribute_dataset(dataset)
+        dist_dataset = dataset
+        # training parameters
+        training_parameters = {
+            # global params
+            'cur_tf_ver': cur_tf_ver,
+            'use_tf_function': args['use_tf_function'],
+            'use_custom_cuda': args['use_custom_cuda'],
+            'model_base_dir': args['model_base_dir'],
+
+            # network params
+            'g_params': g_params,
+            'd_params': d_params,
+
+            # training params
+            'g_opt': {'learning_rate': 0.0002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 8},
+            'd_opt': {'learning_rate': 0.0002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 16},
+            'batch_size': global_batch_size,
+            'n_total_image': 1000,
+            'n_samples': 12,
+            'train_res': args['train_res'],
+        }
+
+        trainer = Trainer(training_parameters, name=f'stylegan2-test-2-{args["train_res"]}x{args["train_res"]}')
+        snapshot1 = tracemalloc.take_snapshot()
+        trainer.train(dist_dataset, strategy)
+        snapshot2 = tracemalloc.take_snapshot()
+
+        top_stats = snapshot2.compare_to(snapshot1, 'lineno')
+
+        print("[ Top 10 differences ]")
+        for stat in top_stats[:100]:
+            print(stat)
+
+        display_top(snapshot2)
+
+    return
+
+
+if __name__ == '__main__':
+    main()

train_old.py

@@ -0,0 +1,445 @@
+import os
+import time
+import argparse
+import wandb
+from cv2 import drawFrameAxes
+import numpy as np
+import tensorflow as tf
+
+from utils import str_to_bool
+from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
+from load_models import load_generator, load_discriminator
+from dataset_ffhq import get_ffhq_dataset
+from dataset import DatasetSimple
+from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_2')
+artifact = wandb.Artifact('code', type='code')
+artifact.add_file('train_old.py')
+run.log_artifact(artifact)
+
+def initiate_models(g_params, d_params, use_custom_cuda):
+    discriminator = load_discriminator(d_params, ckpt_dir=None, custom_cuda=use_custom_cuda)
+    generator = load_generator(g_params=g_params, is_g_clone=False, ckpt_dir=None, custom_cuda=use_custom_cuda)
+    g_clone = load_generator(g_params=g_params, is_g_clone=True, ckpt_dir=None, custom_cuda=use_custom_cuda)
+
+    # set initial g_clone weights same as generator
+    g_clone.set_weights(generator.get_weights())
+    return discriminator, generator, g_clone
+
+
+class Trainer(object):
+    def __init__(self, t_params, name):
+        self.cur_tf_ver = t_params['cur_tf_ver']
+        self.use_tf_function = t_params['use_tf_function']
+        self.use_custom_cuda = t_params['use_custom_cuda']
+        self.model_base_dir = t_params['model_base_dir']
+        self.global_batch_size = t_params['batch_size']
+        self.n_total_image = t_params['n_total_image']
+        self.max_steps = int(np.ceil(self.n_total_image / self.global_batch_size))
+        self.n_samples = min(t_params['batch_size'], t_params['n_samples'])
+        self.train_res = t_params['train_res']
+        self.print_step = 10
+        self.save_step = 1000
+        self.image_summary_step = 10000
+        self.reached_max_steps = False
+        self.log_template = '{:s}, {:s}, {:s}'.format(
+            'step {}: elapsed: {:.2f}s, d_loss: {:.3f}, g_loss: {:.3f}',
+            'd_gan_loss: {:.3f}, g_gan_loss: {:.3f}',
+            'r1_penalty: {:.3f}, pl_penalty: {:.3f}')
+
+        # copy network params
+        self.g_params = t_params['g_params']
+        self.d_params = t_params['d_params']
+
+        # set optimizer params
+        self.global_batch_scaler = 1.0 / float(self.global_batch_size)
+        self.r1_gamma = 10.0
+        self.g_opt = self.update_optimizer_params(t_params['g_opt'])
+        self.d_opt = self.update_optimizer_params(t_params['d_opt'])
+        self.pl_minibatch_shrink = 2
+        self.pl_weight = float(self.pl_minibatch_shrink)
+        self.pl_denorm = tf.math.rsqrt(float(self.train_res) * float(self.train_res))
+        self.pl_decay = 0.01
+        self.pl_mean = tf.Variable(initial_value=0.0, name='pl_mean', trainable=False,
+                                   synchronization=tf.VariableSynchronization.ON_READ,
+                                   aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA)
+
+        # create model: model and optimizer must be created under `strategy.scope`
+        self.discriminator, self.generator, self.g_clone = initiate_models(self.g_params,
+                                                                           self.d_params,
+                                                                           self.use_custom_cuda)
+
+        # set optimizers
+        self.d_optimizer = tf.keras.optimizers.Adam(self.d_opt['learning_rate'],
+                                                    beta_1=self.d_opt['beta1'],
+                                                    beta_2=self.d_opt['beta2'],
+                                                    epsilon=self.d_opt['epsilon'])
+        self.g_optimizer = tf.keras.optimizers.Adam(self.g_opt['learning_rate'],
+                                                    beta_1=self.g_opt['beta1'],
+                                                    beta_2=self.g_opt['beta2'],
+                                                    epsilon=self.g_opt['epsilon'])
+
+        # setup saving locations (object based savings)
+        self.ckpt_dir = os.path.join(self.model_base_dir, name)
+        self.ckpt = tf.train.Checkpoint(d_optimizer=self.d_optimizer,
+                                        g_optimizer=self.g_optimizer,
+                                        discriminator=self.discriminator,
+                                        generator=self.generator,
+                                        g_clone=self.g_clone,
+                                        pl_mean=self.pl_mean)
+        self.manager = tf.train.CheckpointManager(self.ckpt, self.ckpt_dir, max_to_keep=10)
+
+        # try to restore
+        self.ckpt.restore(self.manager.latest_checkpoint)
+        if self.manager.latest_checkpoint:
+            print('Restored from {}'.format(self.manager.latest_checkpoint))
+
+            # check if already trained in this resolution
+            restored_step = self.g_optimizer.iterations.numpy()
+            if restored_step >= self.max_steps:
+                print('Already reached max steps {}/{}'.format(restored_step, self.max_steps))
+                self.reached_max_steps = True
+                return
+        else:
+            print('Not restoring from saved checkpoint')
+
+    @staticmethod
+    def update_optimizer_params(params):
+        params_copy = params.copy()
+        mb_ratio = params_copy['reg_interval'] / (params_copy['reg_interval'] + 1)
+        params_copy['learning_rate'] = params_copy['learning_rate'] * mb_ratio
+        params_copy['beta1'] = params_copy['beta1'] ** mb_ratio
+        params_copy['beta2'] = params_copy['beta2'] ** mb_ratio
+        return params_copy
+
+    def d_train_step(self, dist_inputs):
+        real_images = dist_inputs[0]
+
+        with tf.GradientTape() as d_tape:
+            # compute losses
+            d_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'])
+
+            # scale loss
+            d_loss = tf.reduce_sum(d_loss) * self.global_batch_scaler
+
+        d_gradients = d_tape.gradient(d_loss, self.discriminator.trainable_variables)
+        self.d_optimizer.apply_gradients(zip(d_gradients, self.discriminator.trainable_variables))
+        return d_loss
+
+    def d_train_step_reg(self, dist_inputs):
+        real_images = dist_inputs[0]
+
+        with tf.GradientTape() as d_tape:
+            # compute losses
+            d_gan_loss, r1_penalty = d_logistic_r1_reg(real_images, self.generator, self.discriminator,
+                                                       self.g_params['z_dim'])
+            r1_penalty = r1_penalty * (0.5 * self.r1_gamma) * self.d_opt['reg_interval']
+
+            # scale losses
+            r1_penalty = tf.reduce_sum(r1_penalty) * self.global_batch_scaler
+            d_gan_loss = tf.reduce_sum(d_gan_loss) * self.global_batch_scaler
+
+            # combine
+            d_loss = d_gan_loss + r1_penalty
+
+        d_gradients = d_tape.gradient(d_loss, self.discriminator.trainable_variables)
+        self.d_optimizer.apply_gradients(zip(d_gradients, self.discriminator.trainable_variables))
+        return d_loss, d_gan_loss, r1_penalty
+
+    def g_train_step(self, dist_inputs):
+        real_images = dist_inputs[0]
+
+        with tf.GradientTape() as g_tape:
+            # compute losses
+            g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'])
+
+            # scale loss
+            g_loss = tf.reduce_sum(g_loss) * self.global_batch_scaler
+
+        g_gradients = g_tape.gradient(g_loss, self.generator.trainable_variables)
+        self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_variables))
+        return g_loss
+
+    def g_train_step_reg(self, dist_inputs):
+        real_images = dist_inputs[0]
+
+        with tf.GradientTape() as g_tape:
+            # compute losses
+            g_gan_loss, pl_penalty = g_logistic_ns_pathreg(real_images, self.generator, self.discriminator,
+                                                           self.g_params['z_dim'], self.pl_mean,
+                                                           self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay)
+            pl_penalty = pl_penalty * self.pl_weight * self.g_opt['reg_interval']
+
+            # scale loss
+            pl_penalty = tf.reduce_sum(pl_penalty) * self.global_batch_scaler
+            g_gan_loss = tf.reduce_sum(g_gan_loss) * self.global_batch_scaler
+
+            # combine
+            g_loss = g_gan_loss + pl_penalty
+
+        g_gradients = g_tape.gradient(g_loss, self.generator.trainable_variables)
+        self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_variables))
+        return g_loss, g_gan_loss, pl_penalty
+
+#    @tf.function
+    def train(self, dist_datasets, strategy):
+        def dist_d_train_step(inputs):
+            if self.cur_tf_ver == '2.0.0':
+                per_replica_losses = strategy.experimental_run_v2(fn=self.d_train_step, args=(inputs,))
+            else:
+                per_replica_losses = strategy.run(fn=self.d_train_step, args=(inputs,))
+            mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
+            return mean_d_loss
+
+        def dist_d_train_step_reg(inputs):
+            if self.cur_tf_ver == '2.0.0':
+                per_replica_losses = strategy.experimental_run_v2(fn=self.d_train_step_reg, args=(inputs,))
+            else:
+                per_replica_losses = strategy.run(fn=self.d_train_step_reg, args=(inputs,))
+            mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
+            mean_d_gan_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
+            mean_r1_penalty = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
+            return mean_d_loss, mean_d_gan_loss, mean_r1_penalty
+
+        def dist_g_train_step(inputs):
+            if self.cur_tf_ver == '2.0.0':
+                per_replica_losses = strategy.experimental_run_v2(fn=self.g_train_step, args=(inputs,))
+            else:
+                per_replica_losses = strategy.run(fn=self.g_train_step, args=(inputs,))
+            mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
+            return mean_g_loss
+
+        def dist_g_train_step_reg(inputs):
+            if self.cur_tf_ver == '2.0.0':
+                per_replica_losses = strategy.experimental_run_v2(fn=self.g_train_step_reg, args=(inputs,))
+            else:
+                per_replica_losses = strategy.run(fn=self.g_train_step_reg, args=(inputs,))
+            mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
+            mean_g_gan_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
+            mean_pl_penalty = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
+            return mean_g_loss, mean_g_gan_loss, mean_pl_penalty
+
+        def dist_gen_samples(dist_inputs):
+            if self.cur_tf_ver == '2.0.0':
+                per_replica_samples = strategy.experimental_run_v2(self.gen_samples, args=(dist_inputs,))
+            else:
+                per_replica_samples = strategy.run(self.gen_samples, args=(dist_inputs,))
+            return per_replica_samples
+
+        # wrap with tf.function
+        if self.use_tf_function:
+            dist_d_train_step = tf.function(dist_d_train_step)
+            dist_g_train_step = tf.function(dist_g_train_step)
+            dist_d_train_step_reg = tf.function(dist_d_train_step_reg)
+            dist_g_train_step_reg = tf.function(dist_g_train_step_reg)
+            dist_gen_samples = tf.function(dist_gen_samples)
+
+        if self.reached_max_steps:
+            return
+
+        # start actual training
+        print('Start Training')
+
+        # setup tensorboards
+        train_summary_writer = tf.summary.create_file_writer(self.ckpt_dir)
+
+        # loss metrics
+        metric_d_loss = tf.keras.metrics.Mean('d_loss', dtype=tf.float32)
+        metric_g_loss = tf.keras.metrics.Mean('g_loss', dtype=tf.float32)
+        metric_d_gan_loss = tf.keras.metrics.Mean('d_gan_loss', dtype=tf.float32)
+        metric_g_gan_loss = tf.keras.metrics.Mean('g_gan_loss', dtype=tf.float32)
+        metric_r1_penalty = tf.keras.metrics.Mean('r1_penalty', dtype=tf.float32)
+        metric_pl_penalty = tf.keras.metrics.Mean('pl_penalty', dtype=tf.float32)
+
+        # start training
+        zero = tf.constant(0.0, dtype=tf.float32)
+        print('max_steps: {}'.format(self.max_steps))
+        t_start = time.time()
+        for real_images in dist_datasets:
+            step = self.g_optimizer.iterations.numpy()
+
+            # d train step
+            if (step + 1) % self.d_opt['reg_interval'] == 0:
+                d_loss, d_gan_loss, r1_penalty = dist_d_train_step_reg((real_images, ))
+            else:
+                d_loss = dist_d_train_step((real_images, ))
+                d_gan_loss = d_loss
+                r1_penalty = zero
+
+            # g train step
+            if (step + 1) % self.g_opt['reg_interval'] == 0:
+                g_loss, g_gan_loss, pl_penalty = dist_g_train_step_reg((real_images, ))
+            else:
+                g_loss = dist_g_train_step((real_images,))
+                g_gan_loss = g_loss
+                pl_penalty = zero
+
+            # update g_clone
+            self.g_clone.set_as_moving_average_of(self.generator)
+
+            # update metrics
+            metric_d_loss(d_loss)
+            metric_g_loss(g_loss)
+            metric_d_gan_loss(d_gan_loss)
+            metric_g_gan_loss(g_gan_loss)
+            metric_r1_penalty(r1_penalty)
+            metric_pl_penalty(pl_penalty)
+
+            # get current step
+            step = self.g_optimizer.iterations.numpy()
+
+            # save to tensorboard
+            with train_summary_writer.as_default():
+                tf.summary.scalar('d_loss', metric_d_loss.result(), step=step)
+                tf.summary.scalar('g_loss', metric_g_loss.result(), step=step)
+                tf.summary.scalar('d_gan_loss', metric_d_gan_loss.result(), step=step)
+                tf.summary.scalar('g_gan_loss', metric_g_gan_loss.result(), step=step)
+                tf.summary.scalar('r1_penalty', metric_r1_penalty.result(), step=step)
+                tf.summary.scalar('pl_penalty', metric_pl_penalty.result(), step=step)
+
+            # save every self.save_step
+            if step % self.save_step == 0:
+                self.manager.save(checkpoint_number=step)
+
+            # save every self.image_summary_step
+            if step % self.image_summary_step == 0:
+                # add summary image
+                test_z = tf.random.normal(shape=(self.n_samples, self.g_params['z_dim']), dtype=tf.dtypes.float32)
+                test_labels = tf.ones((self.n_samples, self.g_params['labels_dim']), dtype=tf.dtypes.float32)
+                summary_image = dist_gen_samples((test_z, test_labels))
+
+                # convert to tensor image
+                summary_image = self.convert_per_replica_image(summary_image, strategy)
+
+                with train_summary_writer.as_default():
+                    tf.summary.image('images', summary_image, step=step)
+
+            # print every self.print_steps
+            if step % self.print_step == 0:
+                elapsed = time.time() - t_start
+                print(self.log_template.format(step, elapsed, d_loss.numpy(), g_loss.numpy(),
+                                               d_gan_loss.numpy(), g_gan_loss.numpy(),
+                                               r1_penalty.numpy(), pl_penalty.numpy()))
+
+                # reset timer
+                t_start = time.time()
+
+            # check exit status
+            if step >= self.max_steps:
+                break
+
+        # save last checkpoint
+        step = self.g_optimizer.iterations.numpy()
+        self.manager.save(checkpoint_number=step)
+        return
+
+    def gen_samples(self, inputs):
+        test_z, test_labels = inputs
+
+        # run networks
+        fake_images_05 = self.g_clone([test_z, test_labels], truncation_psi=0.5, training=False)
+        fake_images_07 = self.g_clone([test_z, test_labels], truncation_psi=0.7, training=False)
+
+        # merge on batch dimension: [n_samples, 3, out_res, 2 * out_res]
+        final_image = tf.concat([fake_images_05, fake_images_07], axis=2)
+        return final_image
+
+    @staticmethod
+    def convert_per_replica_image(nchw_per_replica_images, strategy):
+        as_tensor = tf.concat(strategy.experimental_local_results(nchw_per_replica_images), axis=0)
+        as_tensor = tf.transpose(as_tensor, perm=[0, 2, 3, 1])
+        as_tensor = (tf.clip_by_value(as_tensor, -1.0, 1.0) + 1.0) * 127.5
+        as_tensor = tf.cast(as_tensor, tf.uint8)
+        return as_tensor
+
+
+def filter_resolutions_featuremaps(resolutions, featuremaps, res):
+    index = resolutions.index(res)
+    filtered_resolutions = resolutions[:index + 1]
+    filtered_featuremaps = featuremaps[:index + 1]
+    return filtered_resolutions, filtered_featuremaps
+
+
+def main():
+    # global program arguments parser
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--allow_memory_growth', type=str_to_bool, nargs='?', const=True, default=True)
+    parser.add_argument('--debug_split_gpu', type=str_to_bool, nargs='?', const=True, default=False)
+    parser.add_argument('--use_tf_function', type=str_to_bool, nargs='?', const=True, default=True)
+    parser.add_argument('--use_custom_cuda', type=str_to_bool, nargs='?', const=True, default=False)
+    parser.add_argument('--model_base_dir', default='./models', type=str)
+    parser.add_argument('--tfrecord_dir', default='./tfrecords', type=str)
+    parser.add_argument('--train_res', default=512, type=int)
+    parser.add_argument('--shuffle_buffer_size', default=1000, type=int)
+    parser.add_argument('--batch_size_per_replica', default=12, type=int)
+    args = vars(parser.parse_args())
+
+    # check tensorflow version
+    cur_tf_ver = check_tf_version()
+
+    # GPU environment settings
+    if args['allow_memory_growth']:
+        allow_memory_growth()
+    if args['debug_split_gpu']:
+        split_gpu_for_testing(mem_in_gb=4.5)
+
+    # network params
+    resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
+    featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
+    train_resolutions, train_featuremaps = filter_resolutions_featuremaps(resolutions, featuremaps, args['train_res'])
+    g_params = {
+        'z_dim': 512,
+        'w_dim': 512,
+        'labels_dim': 0,
+        'n_mapping': 8,
+        'resolutions': train_resolutions,
+        'featuremaps': train_featuremaps,
+    }
+    d_params = {
+        'labels_dim': 0,
+        'resolutions': train_resolutions,
+        'featuremaps': train_featuremaps,
+    }
+
+    # prepare distribute strategy
+    strategy = tf.distribute.MirroredStrategy()
+    global_batch_size = args['batch_size_per_replica'] * strategy.num_replicas_in_sync
+
+    # prepare dataset
+    # dataset = get_ffhq_dataset(args['tfrecord_dir'], args['train_res'], batch_size=global_batch_size, epochs=None)
+    dataset = DatasetSimple('/home/workspace/local-data/stylegan/npz').shuffle(buffer_size=1000).repeat(None).batch(global_batch_size).prefetch(buffer_size=tf.data.experimental.AUTOTUNE).cache()
+
+    with strategy.scope():
+        # distribute dataset
+        #dist_dataset = strategy.experimental_distribute_dataset(dataset)
+        dist_dataset = dataset
+        # training parameters
+        training_parameters = {
+            # global params
+            'cur_tf_ver': cur_tf_ver,
+            'use_tf_function': args['use_tf_function'],
+            'use_custom_cuda': args['use_custom_cuda'],
+            'model_base_dir': args['model_base_dir'],
+
+            # network params
+            'g_params': g_params,
+            'd_params': d_params,
+
+            # training params
+            'g_opt': {'learning_rate': 0.0002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 8},
+            'd_opt': {'learning_rate': 0.0002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 16},
+            'batch_size': global_batch_size,
+            'n_total_image': 25000000,
+            'n_samples': 12,
+            'train_res': args['train_res'],
+        }
+
+        trainer = Trainer(training_parameters, name=f'stylegan2-test-2-{args["train_res"]}x{args["train_res"]}')
+        trainer.train(dist_dataset, strategy)
+    return
+
+
+if __name__ == '__main__':
+    main()

utils.py

@@ -0,0 +1,8 @@
+def str_to_bool(value):
+    if isinstance(value, bool):
+        return value
+    if value.lower() in {'false', 'f', '0', 'no', 'n'}:
+        return False
+    elif value.lower() in {'true', 't', '1', 'yes', 'y'}:
+        return True
+    raise ValueError(f'{value} is not a valid boolean value')