Fork adriansahlman's stylegan2_pytorch

.gitignore

@@ -0,0 +1,8 @@
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+results/

LICENSE.txt

@@ -0,0 +1,7 @@
+Copyright 2020 Adrian Sahlman
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

requirements.txt

@@ -0,0 +1,9 @@
+numpy
+pillow
+pyyaml
+requests
+scipy
+tensorboard
+torch
+torchvision
+tqdm

run_convert_from_tf.py

@@ -0,0 +1,373 @@
+import os
+import re
+import pickle
+import argparse
+import io
+import requests
+import torch
+import stylegan2
+from stylegan2 import utils
+
+
+pretrained_model_urls = {
+    'car-config-e':                    'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-car-config-e.pkl',
+    'car-config-f':                    'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-car-config-f.pkl',
+    'cat-config-f':                    'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-cat-config-f.pkl',
+    'church-config-f':                 'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-church-config-f.pkl',
+    'ffhq-config-e':                   'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-ffhq-config-e.pkl',
+    'ffhq-config-f':                   'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-ffhq-config-f.pkl',
+    'horse-config-f':                  'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-horse-config-f.pkl',
+    'car-config-e-Gorig-Dorig':        'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-car-config-e-Gorig-Dorig.pkl',
+    'car-config-e-Gorig-Dresnet':      'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-car-config-e-Gorig-Dresnet.pkl',
+    'car-config-e-Gorig-Dskip':        'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-car-config-e-Gorig-Dskip.pkl',
+    'car-config-e-Gresnet-Dorig':      'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-car-config-e-Gresnet-Dorig.pkl',
+    'car-config-e-Gresnet-Dresnet':    'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-car-config-e-Gresnet-Dresnet.pkl',
+    'car-config-e-Gresnet-Dskip':      'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-car-config-e-Gresnet-Dskip.pkl',
+    'car-config-e-Gskip-Dorig':        'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-car-config-e-Gskip-Dorig.pkl',
+    'car-config-e-Gskip-Dresnet':      'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-car-config-e-Gskip-Dresnet.pkl',
+    'car-config-e-Gskip-Dskip':        'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-car-config-e-Gskip-Dskip.pkl',
+    'ffhq-config-e-Gorig-Dorig':       'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-ffhq-config-e-Gorig-Dorig.pkl',
+    'ffhq-config-e-Gorig-Dresnet':     'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-ffhq-config-e-Gorig-Dresnet.pkl',
+    'ffhq-config-e-Gorig-Dskip':       'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-ffhq-config-e-Gorig-Dskip.pkl',
+    'ffhq-config-e-Gresnet-Dorig':     'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-ffhq-config-e-Gresnet-Dorig.pkl',
+    'ffhq-config-e-Gresnet-Dresnet':   'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-ffhq-config-e-Gresnet-Dresnet.pkl',
+    'ffhq-config-e-Gresnet-Dskip':     'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-ffhq-config-e-Gresnet-Dskip.pkl',
+    'ffhq-config-e-Gskip-Dorig':       'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-ffhq-config-e-Gskip-Dorig.pkl',
+    'ffhq-config-e-Gskip-Dresnet':     'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-ffhq-config-e-Gskip-Dresnet.pkl',
+    'ffhq-config-e-Gskip-Dskip':       'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/table2/stylegan2-ffhq-config-e-Gskip-Dskip.pkl',
+}
+
+
+class Unpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        if module == 'dnnlib.tflib.network' and name == 'Network':
+            return utils.AttributeDict
+        return super(Unpickler, self).find_class(module, name)
+
+
+def load_tf_models_file(fpath):
+    with open(fpath, 'rb') as fp:
+        return Unpickler(fp).load()
+
+
+def load_tf_models_url(url):
+    print('Downloading file {}...'.format(url))
+    with requests.Session() as session:
+        with session.get(url) as ret:
+            fp = io.BytesIO(ret.content)
+            return Unpickler(fp).load()
+
+
+def convert_kwargs(static_kwargs, kwargs_mapping):
+    kwargs = utils.AttributeDict()
+    for key, value in static_kwargs.items():
+        if key in kwargs_mapping:
+            if value == 'lrelu':
+                value = 'leaky:0.2'
+            for k in utils.to_list(kwargs_mapping[key]):
+                kwargs[k] = value
+    return kwargs
+
+
+_PERMITTED_MODELS = ['G_main', 'G_mapping', 'G_synthesis_stylegan2', 'D_stylegan2', 'D_main', 'G_synthesis']
+def convert_from_tf(tf_state):
+    tf_state = utils.AttributeDict.convert_dict_recursive(tf_state)
+    model_type = tf_state.build_func_name
+    assert model_type in _PERMITTED_MODELS, \
+        'Found model type {}. '.format(model_type) + \
+        'Allowed model types are: {}'.format(_PERMITTED_MODELS)
+
+    if model_type == 'G_main':
+        kwargs = convert_kwargs(
+            static_kwargs=tf_state.static_kwargs,
+            kwargs_mapping={
+                'dlatent_avg_beta': 'dlatent_avg_beta'
+            }
+        )
+        kwargs.G_mapping = convert_from_tf(tf_state.components.mapping)
+        kwargs.G_synthesis = convert_from_tf(tf_state.components.synthesis)
+        G = stylegan2.models.Generator(**kwargs)
+        for name, var in tf_state.variables:
+            if name == 'dlatent_avg':
+                G.dlatent_avg.data.copy_(torch.from_numpy(var))
+        kwargs = convert_kwargs(
+            static_kwargs=tf_state.static_kwargs,
+            kwargs_mapping={
+                'truncation_psi': 'truncation_psi',
+                'truncation_cutoff': 'truncation_cutoff',
+                'truncation_psi_val': 'truncation_psi',
+                'truncation_cutoff_val': 'truncation_cutoff'
+            }
+        )
+        G.set_truncation(**kwargs)
+        return G
+
+    if model_type == 'G_mapping':
+        kwargs = convert_kwargs(
+            static_kwargs=tf_state.static_kwargs,
+            kwargs_mapping={
+                'mapping_nonlinearity': 'activation',
+                'normalize_latents': 'normalize_input',
+                'mapping_lr_mul': 'lr_mul'
+            }
+        )
+        kwargs.num_layers = sum(
+            1 for var_name, _ in tf_state.variables
+            if re.match('Dense[0-9]+/weight', var_name)
+        )
+        for var_name, var in tf_state.variables:
+            if var_name == 'LabelConcat/weight':
+                kwargs.label_size = var.shape[0]
+            if var_name == 'Dense0/weight':
+                kwargs.latent_size = var.shape[0]
+                kwargs.hidden = var.shape[1]
+            if var_name == 'Dense{}/bias'.format(kwargs.num_layers - 1):
+                kwargs.out_size = var.shape[0]
+        G_mapping = stylegan2.models.GeneratorMapping(**kwargs)
+        for var_name, var in tf_state.variables:
+            if re.match('Dense[0-9]+/[a-zA-Z]*', var_name):
+                layer_idx = int(re.search('Dense(\d+)/[a-zA-Z]*', var_name).groups()[0])
+                if var_name.endswith('weight'):
+                    G_mapping.main[layer_idx].layer.weight.data.copy_(
+                        torch.from_numpy(var.T).contiguous())
+                elif var_name.endswith('bias'):
+                    G_mapping.main[layer_idx].bias.data.copy_(torch.from_numpy(var))
+            if var_name == 'LabelConcat/weight':
+                G_mapping.embedding.weight.data.copy_(torch.from_numpy(var))
+        return G_mapping
+
+    if model_type == 'G_synthesis_stylegan2' or model_type == 'G_synthesis':
+        assert tf_state.static_kwargs.get('fused_modconv', True), \
+            'Can not load TF networks that use `fused_modconv=False`'
+        noise_tensors = []
+        conv_vars = {}
+        for var_name, var in tf_state.variables:
+            if var_name.startswith('noise'):
+                noise_tensors.append(torch.from_numpy(var))
+            else:
+                layer_size = int(re.search('(\d+)x[0-9]+/*', var_name).groups()[0])
+                if layer_size not in conv_vars:
+                    conv_vars[layer_size] = {}
+                var_name = var_name.replace('{}x{}/'.format(layer_size, layer_size), '')
+                conv_vars[layer_size][var_name] = var
+        noise_tensors = sorted(noise_tensors, key=lambda x:x.size(-1))
+        kwargs = convert_kwargs(
+            static_kwargs=tf_state.static_kwargs,
+            kwargs_mapping={
+                'nonlinearity': 'activation',
+                'resample_filter': ['conv_filter', 'skip_filter']
+            }
+        )
+        kwargs.skip = False
+        kwargs.resnet = True
+        kwargs.channels = []
+        for size in sorted(conv_vars.keys(), reverse=True):
+            if size == 4:
+                if 'ToRGB/weight' in conv_vars[size]:
+                    kwargs.skip = True
+                    kwargs.resnet = False
+                kwargs.latent_size = conv_vars[size]['Conv/mod_weight'].shape[0]
+                kwargs.channels.append(conv_vars[size]['Conv/bias'].shape[0])
+            else:
+                kwargs.channels.append(conv_vars[size]['Conv1/bias'].shape[0])
+            if 'ToRGB/bias' in conv_vars[size]:
+                kwargs.data_channels = conv_vars[size]['ToRGB/bias'].shape[0]
+        G_synthesis = stylegan2.models.GeneratorSynthesis(**kwargs)
+        G_synthesis.const.data.copy_(torch.from_numpy(conv_vars[4]['Const/const']).squeeze(0))
+        def assign_weights(layer, weight, bias, mod_weight, mod_bias, noise_strength, transposed=False):
+            layer.bias.data.copy_(torch.from_numpy(bias))
+            layer.layer.weight.data.copy_(torch.tensor(noise_strength))
+            layer.layer.layer.dense.layer.weight.data.copy_(
+                torch.from_numpy(mod_weight.T).contiguous())
+            layer.layer.layer.dense.bias.data.copy_(torch.from_numpy(mod_bias + 1))
+            weight = torch.from_numpy(weight).permute((3, 2, 0, 1)).contiguous()
+            if transposed:
+                weight = weight.flip(dims=[2,3])
+            layer.layer.layer.weight.data.copy_(weight)
+        conv_blocks = G_synthesis.conv_blocks
+        for i, size in enumerate(sorted(conv_vars.keys())):
+            block = conv_blocks[i]
+            if size == 4:
+                assign_weights(
+                    layer=block.conv_block[0],
+                    weight=conv_vars[size]['Conv/weight'],
+                    bias=conv_vars[size]['Conv/bias'],
+                    mod_weight=conv_vars[size]['Conv/mod_weight'],
+                    mod_bias=conv_vars[size]['Conv/mod_bias'],
+                    noise_strength=conv_vars[size]['Conv/noise_strength'],
+                )
+            else:
+                assign_weights(
+                    layer=block.conv_block[0],
+                    weight=conv_vars[size]['Conv0_up/weight'],
+                    bias=conv_vars[size]['Conv0_up/bias'],
+                    mod_weight=conv_vars[size]['Conv0_up/mod_weight'],
+                    mod_bias=conv_vars[size]['Conv0_up/mod_bias'],
+                    noise_strength=conv_vars[size]['Conv0_up/noise_strength'],
+                    transposed=True
+                )
+                assign_weights(
+                    layer=block.conv_block[1],
+                    weight=conv_vars[size]['Conv1/weight'],
+                    bias=conv_vars[size]['Conv1/bias'],
+                    mod_weight=conv_vars[size]['Conv1/mod_weight'],
+                    mod_bias=conv_vars[size]['Conv1/mod_bias'],
+                    noise_strength=conv_vars[size]['Conv1/noise_strength'],
+                )
+                if 'Skip/weight' in conv_vars[size]:
+                    block.projection.weight.data.copy_(torch.from_numpy(
+                        conv_vars[size]['Skip/weight']).permute((3, 2, 0, 1)).contiguous())
+            to_RGB = G_synthesis.to_data_layers[i]
+            if to_RGB is not None:
+                to_RGB.bias.data.copy_(torch.from_numpy(conv_vars[size]['ToRGB/bias']))
+                to_RGB.layer.weight.data.copy_(torch.from_numpy(
+                    conv_vars[size]['ToRGB/weight']).permute((3, 2, 0, 1)).contiguous())
+                to_RGB.layer.dense.bias.data.copy_(
+                    torch.from_numpy(conv_vars[size]['ToRGB/mod_bias'] + 1))
+                to_RGB.layer.dense.layer.weight.data.copy_(
+                    torch.from_numpy(conv_vars[size]['ToRGB/mod_weight'].T).contiguous())
+        if not tf_state.static_kwargs.get('randomize_noise', True):
+            G_synthesis.static_noise(noise_tensors=noise_tensors)
+        return G_synthesis
+
+    if model_type == 'D_stylegan2' or model_type == 'D_main':
+        output_vars = {}
+        conv_vars = {}
+        for var_name, var in tf_state.variables:
+            if var_name.startswith('Output'):
+                output_vars[var_name.replace('Output/', '')] = var
+            else:
+                layer_size = int(re.search('(\d+)x[0-9]+/*', var_name).groups()[0])
+                if layer_size not in conv_vars:
+                    conv_vars[layer_size] = {}
+                var_name = var_name.replace('{}x{}/'.format(layer_size, layer_size), '')
+                conv_vars[layer_size][var_name] = var
+        kwargs = convert_kwargs(
+            static_kwargs=tf_state.static_kwargs,
+            kwargs_mapping={
+                'nonlinearity': 'activation',
+                'resample_filter': ['conv_filter', 'skip_filter'],
+                'mbstd_group_size': 'mbstd_group_size'
+            }
+        )
+        kwargs.skip = False
+        kwargs.resnet = True
+        kwargs.channels = []
+        for size in sorted(conv_vars.keys(), reverse=True):
+            if size == 4:
+                if 'FromRGB/weight' in conv_vars[size]:
+                    kwargs.skip = True
+                    kwargs.resnet = False
+                kwargs.channels.append(conv_vars[size]['Conv/bias'].shape[0])
+                kwargs.dense_hidden = conv_vars[size]['Dense0/bias'].shape[0]
+            else:
+                kwargs.channels.append(conv_vars[size]['Conv0/bias'].shape[0])
+            if 'FromRGB/weight' in conv_vars[size]:
+                kwargs.data_channels = conv_vars[size]['FromRGB/weight'].shape[-2]
+        output_size = output_vars['bias'].shape[0]
+        if output_size > 1:
+            kwargs.label_size = output_size
+        D = stylegan2.models.Discriminator(**kwargs)
+        def assign_weights(layer, weight, bias):
+            layer.bias.data.copy_(torch.from_numpy(bias))
+            layer.layer.weight.data.copy_(
+                torch.from_numpy(weight).permute((3, 2, 0, 1)).contiguous())
+        conv_blocks = D.conv_blocks
+        for i, size in enumerate(sorted(conv_vars.keys())):
+            block = conv_blocks[-i - 1]
+            if size == 4:
+                assign_weights(
+                    layer=block[-1].conv_block[0],
+                    weight=conv_vars[size]['Conv/weight'],
+                    bias=conv_vars[size]['Conv/bias'],
+                )
+            else:
+                assign_weights(
+                    layer=block.conv_block[0],
+                    weight=conv_vars[size]['Conv0/weight'],
+                    bias=conv_vars[size]['Conv0/bias'],
+                )
+                assign_weights(
+                    layer=block.conv_block[1],
+                    weight=conv_vars[size]['Conv1_down/weight'],
+                    bias=conv_vars[size]['Conv1_down/bias'],
+                )
+                if 'Skip/weight' in conv_vars[size]:
+                    block.projection.weight.data.copy_(torch.from_numpy(
+                        conv_vars[size]['Skip/weight']).permute((3, 2, 0, 1)).contiguous())
+            from_RGB = D.from_data_layers[-i - 1]
+            if from_RGB is not None:
+                from_RGB.bias.data.copy_(torch.from_numpy(conv_vars[size]['FromRGB/bias']))
+                from_RGB.layer.weight.data.copy_(torch.from_numpy(
+                    conv_vars[size]['FromRGB/weight']).permute((3, 2, 0, 1)).contiguous())
+        return D
+
+
+def get_arg_parser():
+    parser = argparse.ArgumentParser(
+        description='Convert tensorflow stylegan2 model to pytorch.',
+        epilog='Pretrained models that can be downloaded:\n{}'.format(
+            '\n'.join(pretrained_model_urls.keys()))
+    )
+
+    parser.add_argument(
+        '-i',
+        '--input',
+        help='File path to pickled tensorflow models.',
+        type=str,
+        default=None,
+    )
+
+    parser.add_argument(
+        '-d',
+        '--download',
+        help='Download the specified pretrained model. Use --help for info on available models.',
+        type=str,
+        default=None,
+    )
+
+    parser.add_argument(
+        '-o',
+        '--output',
+        help='One or more output file paths. Alternatively a directory path ' + \
+            'where all models will be saved. Default: current directory',
+        type=str,
+        nargs='*',
+        default=['.'],
+    )
+
+    return parser
+
+
+def main():
+    args = get_arg_parser().parse_args()
+    assert bool(args.input) != bool(args.download), \
+        'Incorrect input format. Can only take either one ' + \
+        'input filepath to a pickled tensorflow model or ' + \
+        'a model name to download, but not both at the same ' + \
+        'time or none at all.'
+    if args.input:
+        unpickled = load_tf_models_file(args.input)
+    else:
+        assert args.download in pretrained_model_urls.keys(), \
+            'Unknown model {}. Use --help for list of models.'.format(args.download)
+        unpickled = load_tf_models_url(pretrained_model_urls[args.download])
+    if not isinstance(unpickled, (tuple, list)):
+        unpickled = [unpickled]
+    print('Converting tensorflow models and saving them...')
+    converted = [convert_from_tf(tf_state) for tf_state in unpickled]
+    if len(args.output) == 1 and (os.path.isdir(args.output[0]) or not os.path.splitext(args.output[0])[-1]):
+        if not os.path.exists(args.output[0]):
+            os.makedirs(args.output[0])
+        for tf_state, torch_model in zip(unpickled, converted):
+            torch_model.save(os.path.join(args.output[0], tf_state['name'] + '.pth'))
+    else:
+        assert len(args.output) == len(converted), 'Found {} models '.format(len(converted)) + \
+            'in pickled file but only {} output paths were given.'.format(len(args.output))
+        for out_path, torch_model in zip(args.output, converted):
+            torch_model.save(out_path)
+    print('Done!')
+
+
+if __name__ == '__main__':
+    main()

run_generator.py

@@ -0,0 +1,361 @@
+import warnings
+import argparse
+import os
+from PIL import Image
+import numpy as np
+import torch
+
+import stylegan2
+from stylegan2 import utils
+
+#----------------------------------------------------------------------------
+
+_description = """StyleGAN2 generator.
+Run 'python %(prog)s <subcommand> --help' for subcommand help."""
+
+#----------------------------------------------------------------------------
+
+_examples = """examples:
+  # Train a network or convert a pretrained one.
+  # Example of converting pretrained ffhq model:
+  python run_convert_from_tf --download ffhq-config-f --output G.pth D.pth Gs.pth
+
+  # Generate ffhq uncurated images (matches paper Figure 12)
+  python %(prog)s generate_images --network=Gs.pth --seeds=6600-6625 --truncation_psi=0.5
+
+  # Generate ffhq curated images (matches paper Figure 11)
+  python %(prog)s generate_images --network=Gs.pth --seeds=66,230,389,1518 --truncation_psi=1.0
+
+  # Example of converting pretrained car model:
+  python run_convert_from_tf --download car-config-f --output G_car.pth D_car.pth Gs_car.pth
+
+  # Generate uncurated car images (matches paper Figure 12)
+  python %(prog)s generate_images --network=Gs_car.pth --seeds=6000-6025 --truncation_psi=0.5
+
+  # Generate style mixing example (matches style mixing video clip)
+  python %(prog)s style_mixing_example --network=Gs.pth --row_seeds=85,100,75,458,1500 --col_seeds=55,821,1789,293 --truncation_psi=1.0
+"""
+
+#----------------------------------------------------------------------------
+
+def _add_shared_arguments(parser):
+
+    parser.add_argument(
+        '--network',
+        help='Network file path',
+        required=True,
+        metavar='FILE'
+    )
+
+    parser.add_argument(
+        '--output',
+        help='Root directory for run results. Default: %(default)s',
+        type=str,
+        default='./results',
+        metavar='DIR'
+    )
+
+    parser.add_argument(
+        '--pixel_min',
+        help='Minumum of the value range of pixels in generated images. ' + \
+            'Default: %(default)s',
+        default=-1,
+        type=float,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--pixel_max',
+        help='Maximum of the value range of pixels in generated images. ' + \
+            'Default: %(default)s',
+        default=1,
+        type=float,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--gpu',
+        help='CUDA device indices (given as separate ' + \
+            'values if multiple, i.e. "--gpu 0 1"). Default: Use CPU',
+        type=int,
+        default=[],
+        nargs='*',
+        metavar='INDEX'
+    )
+
+    parser.add_argument(
+        '--truncation_psi',
+        help='Truncation psi. Default: %(default)s',
+        type=float,
+        default=0.5,
+        metavar='VALUE'
+    )
+
+
+def get_arg_parser():
+    parser = argparse.ArgumentParser(
+        description=_description,
+        epilog=_examples,
+        formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+
+    range_desc = 'NOTE: This is a single argument, where list ' + \
+        'elements are separated by "," and ranges are defined as "a-b". Only integers are allowed.'
+
+    subparsers = parser.add_subparsers(help='Sub-commands', dest='command')
+
+    generate_images_parser = subparsers.add_parser(
+        'generate_images', help='Generate images')
+
+    generate_images_parser.add_argument(
+        '--batch_size',
+        help='Batch size for generator. Default: %(default)s',
+        type=int,
+        default=1,
+        metavar='VALUE'
+    )
+
+    generate_images_parser.add_argument(
+        '--seeds',
+        help='List of random seeds for generating images. ' + range_desc,
+        type=utils.range_type,
+        required=True,
+        metavar='RANGE'
+    )
+
+    _add_shared_arguments(generate_images_parser)
+
+    style_mixing_example_parser = subparsers.add_parser(
+        'style_mixing_example', help='Generate style mixing video')
+
+    style_mixing_example_parser.add_argument(
+        '--row_seeds',
+        help='List of random seeds for image rows. ' + range_desc,
+        type=utils.range_type,
+        required=True,
+        metavar='RANGE'
+    )
+
+    style_mixing_example_parser.add_argument(
+        '--col_seeds',
+        help='List of random seeds for image columns. ' + range_desc,
+        type=utils.range_type,
+        required=True,
+        metavar='RANGE'
+    )
+
+    style_mixing_example_parser.add_argument(
+        '--style_layers',
+        help='Indices of layers to mix style for. ' + \
+            'Default: %(default)s. ' + range_desc,
+        type=utils.range_type,
+        default='0-6',
+        metavar='RANGE'
+    )
+
+    style_mixing_example_parser.add_argument(
+        '--grid',
+        help='Save a grid as well of the style mix. Default: %(default)s',
+        type=utils.bool_type,
+        default=True,
+        const=True,
+        nargs='?',
+        metavar='BOOL'
+    )
+
+    _add_shared_arguments(style_mixing_example_parser)
+
+    return parser
+
+#----------------------------------------------------------------------------
+
+def style_mixing_example(G, args):
+    assert max(args.style_layers) < len(G), \
+        'Style layer indices can not be larger than ' + \
+        'number of style layers ({}) of the generator.'.format(len(G))
+    device = torch.device(args.gpu[0] if args.gpu else 'cpu')
+    if device.index is not None:
+        torch.cuda.set_device(device.index)
+    if len(args.gpu) > 1:
+        warnings.warn('Multi GPU is not available for style mixing example. Using device {}'.format(device))
+    G.to(device)
+    G.static_noise()
+    latent_size, label_size = G.latent_size, G.label_size
+    G_mapping, G_synthesis = G.G_mapping, G.G_synthesis
+
+    all_seeds = list(set(args.row_seeds + args.col_seeds))
+    all_z = torch.stack([torch.from_numpy(np.random.RandomState(seed).randn(latent_size)) for seed in all_seeds])
+    all_z = all_z.to(device=device, dtype=torch.float32)
+    if label_size:
+        labels = torch.zeros(len(all_z), dtype=torch.int64, device=device)
+    else:
+        labels = None
+
+    print('Generating disentangled latents...')
+    with torch.no_grad():
+        all_w = G_mapping(latents=all_z, labels=labels)
+
+    all_w = all_w.unsqueeze(1).repeat(1, len(G_synthesis), 1)
+
+    w_avg = G.dlatent_avg
+
+    if args.truncation_psi != 1:
+        all_w = w_avg + args.truncation_psi * (all_w - w_avg)
+
+    w_dict = {seed: w for seed, w in zip(all_seeds, all_w)}
+
+    all_images = []
+
+    progress = utils.ProgressWriter(len(all_w))
+    progress.write('Generating images...', step=False)
+
+    with torch.no_grad():
+        for w in all_w:
+            all_images.append(G_synthesis(w.unsqueeze(0)))
+            progress.step()
+
+    progress.write('Done!', step=False)
+    progress.close()
+
+    all_images = torch.cat(all_images, dim=0)
+
+    image_dict = {(seed, seed): image for seed, image in zip(all_seeds, all_images)}
+
+    progress = utils.ProgressWriter(len(args.row_seeds) * len(args.col_seeds))
+    progress.write('Generating style-mixed images...', step=False)
+
+    for row_seed in args.row_seeds:
+        for col_seed in args.col_seeds:
+            w = w_dict[row_seed].clone()
+            w[args.style_layers] = w_dict[col_seed][args.style_layers]
+            with torch.no_grad():
+                image_dict[(row_seed, col_seed)] = G_synthesis(w.unsqueeze(0)).squeeze(0)
+            progress.step()
+
+    progress.write('Done!', step=False)
+    progress.close()
+
+    progress = utils.ProgressWriter(len(image_dict))
+    progress.write('Saving images...', step=False)
+
+    for (row_seed, col_seed), image in list(image_dict.items()):
+        image = utils.tensor_to_PIL(
+            image, pixel_min=args.pixel_min, pixel_max=args.pixel_max)
+        image_dict[(row_seed, col_seed)] = image
+        image.save(os.path.join(args.output, '%d-%d.png' % (row_seed, col_seed)))
+        progress.step()
+
+    progress.write('Done!', step=False)
+    progress.close()
+
+    if args.grid:
+        print('\n\nSaving style-mixed grid...')
+        H, W = all_images.size()[2:]
+        canvas = Image.new(
+            'RGB', (W * (len(args.col_seeds) + 1), H * (len(args.row_seeds) + 1)), 'black')
+        for row_idx, row_seed in enumerate([None] + args.row_seeds):
+            for col_idx, col_seed in enumerate([None] + args.col_seeds):
+                if row_seed is None and col_seed is None:
+                    continue
+                key = (row_seed, col_seed)
+                if row_seed is None:
+                    key = (col_seed, col_seed)
+                if col_seed is None:
+                    key = (row_seed, row_seed)
+                canvas.paste(image_dict[key], (W * col_idx, H * row_idx))
+        canvas.save(os.path.join(args.output, 'grid.png'))
+        print('Done!')
+
+#----------------------------------------------------------------------------
+
+def generate_images(G, args):
+    latent_size, label_size = G.latent_size, G.label_size
+    device = torch.device(args.gpu[0] if args.gpu else 'cpu')
+    if device.index is not None:
+        torch.cuda.set_device(device.index)
+    G.to(device)
+    if args.truncation_psi != 1:
+        G.set_truncation(truncation_psi=args.truncation_psi)
+    if len(args.gpu) > 1:
+        warnings.warn(
+            'Noise can not be randomized based on the seed ' + \
+            'when using more than 1 GPU device. Noise will ' + \
+            'now be randomized from default random state.'
+        )
+        G.random_noise()
+        G = torch.nn.DataParallel(G, device_ids=args.gpu)
+    else:
+        noise_reference = G.static_noise()
+
+    def get_batch(seeds):
+        latents = []
+        labels = []
+        if len(args.gpu) <= 1:
+            noise_tensors = [[] for _ in noise_reference]
+        for seed in seeds:
+            rnd = np.random.RandomState(seed)
+            latents.append(torch.from_numpy(rnd.randn(latent_size)))
+            if len(args.gpu) <= 1:
+                for i, ref in enumerate(noise_reference):
+                    noise_tensors[i].append(torch.from_numpy(rnd.randn(*ref.size()[1:])))
+            if label_size:
+                labels.append(torch.tensor([rnd.randint(0, label_size)]))
+        latents = torch.stack(latents, dim=0).to(device=device, dtype=torch.float32)
+        if labels:
+            labels = torch.cat(labels, dim=0).to(device=device, dtype=torch.int64)
+        else:
+            labels = None
+        if len(args.gpu) <= 1:
+            noise_tensors = [
+                torch.stack(noise, dim=0).to(device=device, dtype=torch.float32)
+                for noise in noise_tensors
+            ]
+        else:
+            noise_tensors = None
+        return latents, labels, noise_tensors
+
+    progress = utils.ProgressWriter(len(args.seeds))
+    progress.write('Generating images...', step=False)
+
+    for i in range(0, len(args.seeds), args.batch_size):
+        latents, labels, noise_tensors = get_batch(args.seeds[i: i + args.batch_size])
+        if noise_tensors is not None:
+            G.static_noise(noise_tensors=noise_tensors)
+        with torch.no_grad():
+            generated = G(latents, labels=labels)
+        images = utils.tensor_to_PIL(
+            generated, pixel_min=args.pixel_min, pixel_max=args.pixel_max)
+        for seed, img in zip(args.seeds[i: i + args.batch_size], images):
+            img.save(os.path.join(args.output, 'seed%04d.png' % seed))
+            progress.step()
+
+    progress.write('Done!', step=False)
+    progress.close()
+
+#----------------------------------------------------------------------------
+
+def main():
+    args = get_arg_parser().parse_args()
+    assert args.command, 'Missing subcommand.'
+    assert os.path.isdir(args.output) or not os.path.splitext(args.output)[-1], \
+        '--output argument should specify a directory, not a file.'
+    if not os.path.exists(args.output):
+        os.makedirs(args.output)
+
+    G = stylegan2.models.load(args.network)
+    G.eval()
+
+    assert isinstance(G, stylegan2.models.Generator), 'Model type has to be ' + \
+        'stylegan2.models.Generator. Found {}.'.format(type(G))
+
+    if args.command == 'generate_images':
+        generate_images(G, args)
+    elif args.command == 'style_mixing_example':
+        style_mixing_example(G, args)
+    else:
+        raise TypeError('Unkown command {}'.format(args.command))
+
+#----------------------------------------------------------------------------
+
+if __name__ == '__main__':
+    main()

run_metrics.py

@@ -0,0 +1,338 @@
+import os
+import json
+import argparse
+import numpy as np
+import torch
+
+import stylegan2
+from stylegan2 import utils
+
+#----------------------------------------------------------------------------
+
+_description = """Metrics evaluation.
+Run 'python %(prog)s <subcommand> --help' for subcommand help."""
+
+#----------------------------------------------------------------------------
+
+_examples = """examples:
+  # Train a network or convert a pretrained one. In this example we first convert a pretrained one.
+  python run_convert_from_tf --download ffhq-config-f --output G.pth D.pth Gs.pth
+
+  # Project generated images
+  python %(prog)s project_generated_images --network=Gs.pth --seeds=0,1,5
+
+  # Project real images
+  python %(prog)s project_real_images --network=Gs.pth --data-dir=path/to/image_folder
+"""
+
+#----------------------------------------------------------------------------
+
+def _add_shared_arguments(parser):
+
+    parser.add_argument(
+        '--network',
+        help='Network file path',
+        required=True,
+        metavar='FILE'
+    )
+
+    parser.add_argument(
+        '--num_samples',
+        type=int,
+        help='Number of samples to gather for evaluating ' + \
+            'this metric. Default: %(default)s',
+        default=50000,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--size',
+        type=int,
+        help='Rescale images so that this is the size of their ' + \
+            'smallest side in pixels. Default: Unscaled',
+        default=None,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--batch_size',
+        help='Batch size for generator. Default: %(default)s',
+        type=int,
+        default=1,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--output',
+        help='Root directory for run results. Default: %(default)s',
+        type=str,
+        default='./results',
+        metavar='DIR'
+    )
+
+    parser.add_argument(
+        '--pixel_min',
+        help='Minumum of the value range of pixels in generated images. ' + \
+            'Default: %(default)s',
+        default=-1,
+        type=float,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--pixel_max',
+        help='Maximum of the value range of pixels in generated images. ' + \
+            'Default: %(default)s',
+        default=1,
+        type=float,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--gpu',
+        help='CUDA device indices (given as separate ' + \
+            'values if multiple, i.e. "--gpu 0 1"). Default: Use CPU',
+        type=int,
+        default=[],
+        nargs='*',
+        metavar='INDEX'
+    )
+
+
+def get_arg_parser():
+    parser = argparse.ArgumentParser(
+        description=_description,
+        epilog=_examples,
+        formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+
+    subparsers = parser.add_subparsers(help='Sub-commands', dest='command')
+
+    fid_parser = subparsers.add_parser('fid', help='Calculate FID')
+
+    fid_parser.add_argument(
+        '--data_dir',
+        help='Dataset root directory',
+        required=True,
+        metavar='DIR'
+    )
+
+    fid_parser.add_argument(
+        '--reals_batch_size',
+        help='Batch size for gathering statistics of reals. Default: %(default)s',
+        type=int,
+        default=1,
+        metavar='VALUE'
+    )
+
+    fid_parser.add_argument(
+        '--reals_data_workers',
+        help='Data workers for fetching real data samples. Default: %(default)s',
+        type=int,
+        default=4,
+        metavar='VALUE'
+    )
+
+    fid_parser.add_argument(
+        '--truncation_psi',
+        help='Truncation psi. Default: %(default)s',
+        type=float,
+        default=1.0,
+        metavar='VALUE'
+    )
+
+    _add_shared_arguments(fid_parser)
+
+    ppl_parser = subparsers.add_parser('ppl', help='Calculate PPL')
+
+    ppl_parser.add_argument(
+        '--epsilon',
+        type=float,
+        help='Perturbation value. Default: %(default)s',
+        default=1e-4,
+        metavar='VALUE'
+    )
+
+    ppl_parser.add_argument(
+        '--use_dlatent',
+        type=utils.bool_type,
+        help='Measure on perturbations of disentangled latents ' + \
+            'instead of raw latents. Default: %(default)s',
+        default=True,
+        const=True,
+        nargs='?',
+        metavar='BOOL'
+    )
+
+    ppl_parser.add_argument(
+        '--full_sampling',
+        type=utils.bool_type,
+        help='Measure on random interpolation between two inputs ' + \
+            'instead of directly on one input. Default: %(default)s',
+        default=False,
+        const=True,
+        nargs='?',
+        metavar='BOOL'
+    )
+
+    parser.add_argument(
+        '--ppl_ffhq_crop',
+        help='Crop images evaluated for PPL with crop values ' + \
+            'for FFHQ. Default: False',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=False,
+        metavar='BOOL'
+    )
+
+    _add_shared_arguments(ppl_parser)
+
+    return parser
+
+#----------------------------------------------------------------------------
+
+def _report_metric(value, name, args):
+    fpath = os.path.join(args.output, 'metrics.json')
+    metrics = {}
+    if os.path.exists(fpath):
+        with open(fpath, 'r') as fp:
+            try:
+                metrics = json.load(fp)
+            except Exception:
+                pass
+    metrics[name] = value
+    with open(fpath, 'w') as fp:
+        json.dump(metrics, fp)
+    print('\n\nMetric evaluated!:')
+    print('{}: {}'.format(name, value))
+
+#----------------------------------------------------------------------------
+
+def eval_fid(G, prior_generator, args):
+    assert args.data_dir, '--data_dir has to be specified.'
+    dataset = utils.ImageFolder(
+        args.data_dir,
+        pixel_min=args.pixel_min,
+        pixel_max=args.pixel_max
+    )
+    assert len(dataset), 'No images found at {}'.format(args.data_dir)
+
+    inception = stylegan2.external_models.inception.InceptionV3FeatureExtractor(
+        pixel_min=args.pixel_min, pixel_max=args.pixel_max)
+
+    if len(args.gpu) > 1:
+        inception = torch.nn.DataParallel(inception, device_ids=args.gpu)
+
+    args.reals_batch_size = max(args.reals_batch_size, len(args.gpu))
+
+    fid = stylegan2.metrics.fid.FID(
+        G=G,
+        prior_generator=prior_generator,
+        dataset=dataset,
+        num_samples=args.num_samples,
+        fid_model=inception,
+        fid_size=args.size,
+        truncation_psi=args.truncation_psi,
+        reals_batch_size=args.reals_batch_size,
+        reals_data_workers=args.reals_data_workers
+    )
+
+    value = fid.evaluate()
+
+    name = 'FID'
+    if args.size:
+        name += '({})'.format(args.size)
+    if args.truncation_psi != 1:
+        name +='trunc{}'.format(args.truncation_psi)
+    name += ':{}k'.format(args.num_samples // 1000)
+
+    _report_metric(value, name, args)
+
+#----------------------------------------------------------------------------
+
+def eval_ppl(G, prior_generator, args):
+
+    lpips = stylegan2.external_models.lpips.LPIPS_VGG16(
+        pixel_min=args.pixel_min, pixel_max=args.pixel_max)
+
+    if len(args.gpu) > 1:
+        lpips = torch.nn.DataParallel(lpips, device_ids=args.gpu)
+
+    crop = None
+    if args.ppl_ffhq_crop:
+        crop = stylegan2.metrics.ppl.PPL.FFHQ_CROP
+
+    ppl = stylegan2.metrics.ppl.PPL(
+        G=G,
+        prior_generator=prior_generator,
+        num_samples=args.num_samples,
+        epsilon=args.epsilon,
+        use_dlatent=args.use_dlatent,
+        full_sampling=args.full_sampling,
+        crop=crop,
+        lpips_model=lpips,
+        lpips_size=args.size,
+    )
+
+    value = ppl.evaluate()
+
+    name = 'PPL'
+    if args.size:
+        name += '({})'.format(args.size)
+    if args.use_dlatent:
+        name += 'W'
+    else:
+        name += 'Z'
+    if args.full_sampling:
+        name += '-full'
+    else:
+        name += '-end'
+    name += ':{}k'.format(args.num_samples // 1000)
+
+    _report_metric(value, name, args)
+
+#----------------------------------------------------------------------------
+
+def main():
+    args = get_arg_parser().parse_args()
+    assert args.command, 'Missing subcommand.'
+    assert os.path.isdir(args.output) or not os.path.splitext(args.output)[-1], \
+        '--output argument should specify a directory, not a file.'
+    if not os.path.exists(args.output):
+        os.makedirs(args.output)
+
+    G = stylegan2.models.load(args.network)
+    assert isinstance(G, stylegan2.models.Generator), 'Model type has to be ' + \
+        'stylegan2.models.Generator. Found {}.'.format(type(G))
+
+    latent_size, label_size = G.latent_size, G.label_size
+
+    device = torch.device(args.gpu[0] if args.gpu else 'cpu')
+    if device.index is not None:
+        torch.cuda.set_device(device.index)
+
+    G.to(device).eval().requires_grad_(False)
+
+    if len(args.gpu) > 1:
+        G = torch.nn.DataParallel(G, device_ids=args.gpu)
+
+    args.batch_size = max(args.batch_size, len(args.gpu))
+
+    prior_generator = utils.PriorGenerator(
+        latent_size=latent_size,
+        label_size=label_size,
+        batch_size=args.batch_size,
+        device=device
+    )
+
+    if args.command == 'fid':
+        eval_fid(G, prior_generator, args)
+    elif args.command == 'ppl':
+        eval_ppl(G, prior_generator, args)
+    else:
+        raise TypeError('Unkown command {}'.format(args.command))
+
+
+if __name__ == '__main__':
+    main()

run_projector.py

@@ -0,0 +1,401 @@
+import os
+import argparse
+import numpy as np
+import torch
+
+import stylegan2
+from stylegan2 import utils
+
+#----------------------------------------------------------------------------
+
+_description = """StyleGAN2 projector.
+Run 'python %(prog)s <subcommand> --help' for subcommand help."""
+
+#----------------------------------------------------------------------------
+
+_examples = """examples:
+  # Train a network or convert a pretrained one.
+  # Example of converting pretrained ffhq model:
+  python run_convert_from_tf --download ffhq-config-f --output G.pth D.pth Gs.pth
+
+  # Project generated images
+  python %(prog)s project_generated_images --network=Gs.pth --seeds=0,1,5
+
+  # Project real images
+  python %(prog)s project_real_images --network=Gs.pth --data-dir=path/to/image_folder
+
+"""
+
+#----------------------------------------------------------------------------
+
+def _add_shared_arguments(parser):
+
+    parser.add_argument(
+        '--network',
+        help='Network file path',
+        required=True,
+        metavar='FILE'
+    )
+
+    parser.add_argument(
+        '--num_steps',
+        type=int,
+        help='Number of steps to use for projection. ' + \
+            'Default: %(default)s',
+        default=1000,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--batch_size',
+        help='Batch size. Default: %(default)s',
+        type=int,
+        default=1,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--label',
+        help='Label to use for dlatent statistics gathering ' + \
+            '(should be integer index of class). Default: no label.',
+        type=int,
+        default=None,
+        metavar='CLASS_INDEX'
+    )
+
+    parser.add_argument(
+        '--initial_learning_rate',
+        help='Initial learning rate of projection. Default: %(default)s',
+        default=0.1,
+        type=float,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--initial_noise_factor',
+        help='Initial noise factor of projection. Default: %(default)s',
+        default=0.05,
+        type=float,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--lr_rampdown_length',
+        help='Learning rate rampdown length for projection. ' + \
+            'Should be in range [0, 1]. Default: %(default)s',
+        default=0.25,
+        type=float,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--lr_rampup_length',
+        help='Learning rate rampup length for projection. ' + \
+            'Should be in range [0, 1]. Default: %(default)s',
+        default=0.05,
+        type=float,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--noise_ramp_length',
+        help='Learning rate rampdown length for projection. ' + \
+            'Should be in range [0, 1]. Default: %(default)s',
+        default=0.75,
+        type=float,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--regularize_noise_weight',
+        help='The weight for noise regularization. Default: %(default)s',
+        default=1e5,
+        type=float,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--output',
+        help='Root directory for run results. Default: %(default)s',
+        type=str,
+        default='./results',
+        metavar='DIR'
+    )
+
+    parser.add_argument(
+        '--num_snapshots',
+        help='Number of snapshots. Default: %(default)s',
+        type=int,
+        default=5,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--pixel_min',
+        help='Minumum of the value range of pixels in generated images. ' + \
+            'Default: %(default)s',
+        default=-1,
+        type=float,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--pixel_max',
+        help='Maximum of the value range of pixels in generated images. ' + \
+            'Default: %(default)s',
+        default=1,
+        type=float,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--gpu',
+        help='CUDA device indices (given as separate ' + \
+            'values if multiple, i.e. "--gpu 0 1"). Default: Use CPU',
+        type=int,
+        default=[],
+        nargs='*',
+        metavar='INDEX'
+    )
+
+#----------------------------------------------------------------------------
+
+def get_arg_parser():
+    parser = argparse.ArgumentParser(
+        description=_description,
+        epilog=_examples,
+        formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+
+    range_desc = 'NOTE: This is a single argument, where list ' + \
+        'elements are separated by "," and ranges are defined as "a-b". ' + \
+        'Only integers are allowed.'
+
+    subparsers = parser.add_subparsers(help='Sub-commands', dest='command')
+
+    project_generated_images_parser = subparsers.add_parser(
+        'project_generated_images', help='Project generated images')
+
+    project_generated_images_parser.add_argument(
+        '--seeds',
+        help='List of random seeds for generating images. ' + \
+            'Default: 66,230,389,1518. ' + range_desc,
+        type=utils.range_type,
+        default=[66, 230, 389, 1518],
+        metavar='RANGE'
+    )
+
+    project_generated_images_parser.add_argument(
+        '--truncation_psi',
+        help='Truncation psi. Default: %(default)s',
+        type=float,
+        default=1.0,
+        metavar='VALUE'
+    )
+
+    _add_shared_arguments(project_generated_images_parser)
+
+    project_real_images_parser = subparsers.add_parser(
+        'project_real_images', help='Project real images')
+
+    project_real_images_parser.add_argument(
+        '--data_dir',
+        help='Dataset root directory',
+        type=str,
+        required=True,
+        metavar='DIR'
+    )
+
+    project_real_images_parser.add_argument(
+        '--seed',
+        help='When there are more images available than ' + \
+            'the number that is going to be projected this ' + \
+            'seed is used for picking samples. Default: %(default)s',
+        type=int,
+        default=1234,
+        metavar='VALUE'
+    )
+
+    project_real_images_parser.add_argument(
+        '--num_images',
+        type=int,
+        help='Number of images to project. Default: %(default)s',
+        default=3,
+        metavar='VALUE'
+    )
+
+    _add_shared_arguments(project_real_images_parser)
+
+    return parser
+
+#----------------------------------------------------------------------------
+
+def project_images(G, images, name_prefix, args):
+
+    device = torch.device(args.gpu[0] if args.gpu else 'cpu')
+    if device.index is not None:
+        torch.cuda.set_device(device.index)
+    if len(args.gpu) > 1:
+        warnings.warn(
+            'Multi GPU is not available for projection. ' + \
+            'Using device {}'.format(device)
+        )
+    G = utils.unwrap_module(G).to(device)
+
+    lpips_model = stylegan2.external_models.lpips.LPIPS_VGG16(
+        pixel_min=args.pixel_min, pixel_max=args.pixel_max)
+
+    proj = stylegan2.project.Projector(
+        G=G,
+        dlatent_avg_samples=10000,
+        dlatent_avg_label=args.label,
+        dlatent_device=device,
+        dlatent_batch_size=1024,
+        lpips_model=lpips_model,
+        lpips_size=256
+    )
+
+    for i in range(0, len(images), args.batch_size):
+        target = images[i: i + args.batch_size]
+        proj.start(
+            target=target,
+            num_steps=args.num_steps,
+            initial_learning_rate=args.initial_learning_rate,
+            initial_noise_factor=args.initial_noise_factor,
+            lr_rampdown_length=args.lr_rampdown_length,
+            lr_rampup_length=args.lr_rampup_length,
+            noise_ramp_length=args.noise_ramp_length,
+            regularize_noise_weight=args.regularize_noise_weight,
+            verbose=True,
+            verbose_prefix='Projecting image(s) {}/{}'.format(
+                i * args.batch_size + len(target), len(images))
+        )
+        snapshot_steps = set(
+            args.num_steps - np.linspace(
+                0, args.num_steps, args.num_snapshots, endpoint=False, dtype=int))
+        for k, image in enumerate(
+        utils.tensor_to_PIL(target, pixel_min=args.pixel_min, pixel_max=args.pixel_max)):
+            image.save(os.path.join(args.output, name_prefix[i + k] + 'target.png'))
+        for j in range(args.num_steps):
+            proj.step()
+            if j in snapshot_steps:
+                generated = utils.tensor_to_PIL(
+                    proj.generate(), pixel_min=args.pixel_min, pixel_max=args.pixel_max)
+                for k, image in enumerate(generated):
+                    image.save(os.path.join(
+                        args.output, name_prefix[i + k] + 'step%04d.png' % (j + 1)))
+
+#----------------------------------------------------------------------------
+
+def project_generated_images(G, args):
+    latent_size, label_size = G.latent_size, G.label_size
+    device = torch.device(args.gpu[0] if args.gpu else 'cpu')
+    if device.index is not None:
+        torch.cuda.set_device(device.index)
+    G.to(device)
+    if len(args.gpu) > 1:
+        warnings.warn(
+            'Noise can not be randomized based on the seed ' + \
+            'when using more than 1 GPU device. Noise will ' + \
+            'now be randomized from default random state.'
+        )
+        G.random_noise()
+        G = torch.nn.DataParallel(G, device_ids=args.gpu)
+    else:
+        noise_reference = G.static_noise()
+
+    def get_batch(seeds):
+        latents = []
+        labels = []
+        if len(args.gpu) <= 1:
+            noise_tensors = [[] for _ in noise_reference]
+        for seed in seeds:
+            rnd = np.random.RandomState(seed)
+            latents.append(torch.from_numpy(rnd.randn(latent_size)))
+            if len(args.gpu) <= 1:
+                for i, ref in enumerate(noise_reference):
+                    noise_tensors[i].append(
+                        torch.from_numpy(rnd.randn(*ref.size()[1:])))
+            if label_size:
+                labels.append(torch.tensor([rnd.randint(0, label_size)]))
+        latents = torch.stack(latents, dim=0).to(device=device, dtype=torch.float32)
+        if labels:
+            labels = torch.cat(labels, dim=0).to(device=device, dtype=torch.int64)
+        else:
+            labels = None
+        if len(args.gpu) <= 1:
+            noise_tensors = [
+                torch.stack(noise, dim=0).to(device=device, dtype=torch.float32)
+                for noise in noise_tensors
+            ]
+        else:
+            noise_tensors = None
+        return latents, labels, noise_tensors
+
+    images = []
+
+    progress = utils.ProgressWriter(len(args.seeds))
+    progress.write('Generating images...', step=False)
+
+    for i in range(0, len(args.seeds), args.batch_size):
+        latents, labels, noise_tensors = get_batch(args.seeds[i: i + args.batch_size])
+        if noise_tensors is not None:
+            G.static_noise(noise_tensors=noise_tensors)
+        with torch.no_grad():
+            images.append(G(latents, labels=labels))
+        progress.step()
+
+    images = torch.cat(images, dim=0)
+
+    progress.write('Done!', step=False)
+    progress.close()
+
+    name_prefix = ['seed%04d-' % seed for seed in args.seeds]
+    project_images(G, images, name_prefix, args)
+
+#----------------------------------------------------------------------------
+
+def project_real_images(G, args):
+    device = torch.device(args.gpu[0] if args.gpu else 'cpu')
+    print('Loading images from "%s"...' % args.data_dir)
+    dataset = utils.ImageFolder(
+        args.data_dir, pixel_min=args.pixel_min, pixel_max=args.pixel_max)
+
+    rnd = np.random.RandomState(args.seed)
+    indices = rnd.choice(
+        len(dataset), size=min(args.num_images, len(dataset)), replace=False)
+    images = []
+    for i in indices:
+        data = dataset[i]
+        if isinstance(data, (tuple, list)):
+            data = data[0]
+        images.append(data)
+    images = torch.stack(images).to(device)
+    name_prefix = ['image%04d-' % i for i in indices]
+    print('Done!')
+    project_images(G, images, name_prefix, args)
+
+#----------------------------------------------------------------------------
+
+def main():
+    args = get_arg_parser().parse_args()
+    assert args.command, 'Missing subcommand.'
+    assert os.path.isdir(args.output) or not os.path.splitext(args.output)[-1], \
+        '--output argument should specify a directory, not a file.'
+    if not os.path.exists(args.output):
+        os.makedirs(args.output)
+
+    G = stylegan2.models.load(args.network)
+    assert isinstance(G, stylegan2.models.Generator), 'Model type has to be ' + \
+        'stylegan2.models.Generator. Found {}.'.format(type(G))
+
+    if args.command == 'project_generated_images':
+        project_generated_images(G, args)
+    elif args.command == 'project_real_images':
+        project_real_images(G, args)
+    else:
+        raise TypeError('Unkown command {}'.format(args.command))
+
+
+if __name__ == '__main__':
+    main()

run_training.py

@@ -0,0 +1,1002 @@
+import warnings
+import os
+import torch
+from torch import multiprocessing as mp
+import stylegan2
+from stylegan2 import utils
+from stylegan2.external_models import inception, lpips
+from stylegan2.metrics import fid, ppl
+
+#----------------------------------------------------------------------------
+
+def get_arg_parser():
+    parser = utils.ConfigArgumentParser()
+
+    parser.add_argument(
+        '--output',
+        help='Output directory for model weights.',
+        type=str,
+        default=None,
+        metavar='DIR'
+    )
+
+    #----------------------------------------------------------------------------
+    # Model options
+
+    parser.add_argument(
+        '--channels',
+        help='Specify the channels for each layer (can be overriden for individual ' + \
+            'networks with "--g_channels" and "--d_channels". ' + \
+            'Default: %(default)s',
+        nargs='*',
+        type=int,
+        default=[32, 32, 64, 128, 256, 512, 512, 512, 512],
+        metavar='CHANNELS'
+    )
+
+    parser.add_argument(
+        '--latent',
+        help='Size of the prior (noise vector). Default: %(default)s',
+        type=int,
+        default=512,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--label',
+        help='Number of unique labels. Unused if not specified.',
+        type=int,
+        default=0,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--base_shape',
+        help='Data shape of first layer in generator or ' + \
+            'last layer in discriminator. Default: %(default)s',
+        nargs=2,
+        type=int,
+        default=(4, 4),
+        metavar='SIZE'
+    )
+
+    parser.add_argument(
+        '--kernel_size',
+        help='Size of conv kernel. Default: %(default)s',
+        type=int,
+        default=3,
+        metavar='SIZE'
+    )
+
+    parser.add_argument(
+        '--pad_once',
+        help='Pad filtered convs only once before filter instead ' + \
+            'of twice. Default: %(default)s',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=True,
+        metavar='BOOL'
+    )
+
+    parser.add_argument(
+        '--pad_mode',
+        help='Padding mode for conv layers. Default: %(default)s',
+        type=str,
+        default='constant',
+        metavar='MODE'
+    )
+
+    parser.add_argument(
+        '--pad_constant',
+        help='Padding constant for conv layers when `pad_mode` is ' + \
+            '\'constant\'. Default: %(default)s',
+        type=float,
+        default=0,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--filter_pad_mode',
+        help='Padding mode for filter layers. Default: %(default)s',
+        type=str,
+        default='constant',
+        metavar='MODE'
+    )
+
+    parser.add_argument(
+        '--filter_pad_constant',
+        help='Padding constant for filter layers when `filter_pad_mode` ' + \
+            'is \'constant\'. Default: %(default)s',
+        type=float,
+        default=0,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--filter',
+        help='Filter to use whenever FIR is applied. Default: %(default)s',
+        nargs='*',
+        type=float,
+        default=[1, 3, 3, 1],
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--weight_scale',
+        help='Use weight scaling for equalized learning rate. Default: %(default)s',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=True,
+        metavar='BOOL'
+    )
+
+    #----------------------------------------------------------------------------
+    # Generator options
+
+    parser.add_argument(
+        '--g_file',
+        help='Load a generator model from a file instead of constructing a new one. Disabled unless a file is specified.',
+        type=str,
+        default=None,
+        metavar='FILE'
+    )
+
+    parser.add_argument(
+        '--g_channels',
+        help='Instead of the values of "--channels", ' + \
+            'use these for the generator instead.',
+        nargs='*',
+        type=int,
+        default=[],
+        metavar='CHANNELS'
+    )
+
+    parser.add_argument(
+        '--g_skip',
+        help='Use skip connections for the generator. Default: %(default)s',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=True,
+        metavar='BOOL'
+    )
+
+    parser.add_argument(
+        '--g_resnet',
+        help='Use resnet connections for the generator. Default: %(default)s',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=False,
+        metavar='BOOL'
+    )
+
+    parser.add_argument(
+        '--g_conv_block_size',
+        help='Number of layers in a conv block in the generator. Default: %(default)s',
+        type=int,
+        default=2,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--g_normalize',
+        help='Normalize conv features for generator. Default: %(default)s',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=True,
+        metavar='BOOL'
+    )
+
+    parser.add_argument(
+        '--g_fused_conv',
+        help='Fuse conv & upsample into a transposed ' + \
+            'conv for the generator. Default: %(default)s',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=True,
+        metavar='BOOL'
+    )
+
+    parser.add_argument(
+        '--g_activation',
+        help='The non-linear activaiton function for ' + \
+            'the generator. Default: %(default)s',
+        default='leaky:0.2',
+        type=str,
+        metavar='ACTIVATION'
+    )
+
+    parser.add_argument(
+        '--g_conv_resample_mode',
+        help='Resample mode for upsampling conv ' + \
+            'layers for generator. Default: %(default)s',
+        type=str,
+        default='FIR',
+        metavar='MODE'
+    )
+
+    parser.add_argument(
+        '--g_skip_resample_mode',
+        help='Resample mode for skip connection ' + \
+            'upsamples for the generator. Default: %(default)s',
+        type=str,
+        default='FIR',
+        metavar='MODE'
+    )
+
+    parser.add_argument(
+        '--g_lr',
+        help='The learning rate for the generator. Default: %(default)s',
+        default=2e-3,
+        type=float,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--g_betas',
+        help='Beta values for the generator Adam optimizer. Default: %(default)s',
+        type=float,
+        nargs=2,
+        default=(0, 0.99),
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--g_loss',
+        help='Loss function for the generator. Default: %(default)s',
+        default='logistic_ns',
+        type=str,
+        metavar='LOSS'
+    )
+
+    parser.add_argument(
+        '--g_reg',
+        help='Regularization function for the generator with an optional weight (:?). Default: %(default)s',
+        default='pathreg:2',
+        type=str,
+        metavar='REG'
+    )
+
+    parser.add_argument(
+        '--g_reg_interval',
+        help='Interval at which to regularize the generator. Default: %(default)s',
+        default=4,
+        type=int,
+        metavar='INTERVAL'
+    )
+
+    parser.add_argument(
+        '--g_iter',
+        help='Number of generator iterations per training iteration. Default: %(default)s',
+        default=1,
+        type=int,
+        metavar='ITER'
+    )
+
+    parser.add_argument(
+        '--style_mix',
+        help='The probability of passing more than one ' + \
+            'latent to the generator. Default: %(default)s',
+        type=float,
+        default=0.9,
+        metavar='PROBABILITY'
+    )
+
+    parser.add_argument(
+        '--latent_mapping_layers',
+        help='The number of layers of the latent mapping network. Default: %(default)s',
+        type=int,
+        default=8,
+        metavar='LAYERS'
+    )
+
+    parser.add_argument(
+        '--latent_mapping_lr_mul',
+        help='The learning rate multiplier for the latent ' + \
+            'mapping network. Default: %(default)s',
+        type=float,
+        default=0.01,
+        metavar='LR_MUL'
+    )
+
+    parser.add_argument(
+        '--normalize_latent',
+        help='Normalize latent inputs. Default: %(default)s',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=True,
+        metavar='BOOL'
+    )
+
+    parser.add_argument(
+        '--modulate_rgb',
+        help='Modulate RGB layers (use style for output ' + \
+            'layers of generator). Default: %(default)s',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=True,
+        metavar='BOOL'
+    )
+
+    #----------------------------------------------------------------------------
+    # Discriminator options
+
+    parser.add_argument(
+        '--d_file',
+        help='Load a discriminator model from a file instead of constructing a new one. Disabled unless a file is specified.',
+        type=str,
+        default=None,
+        metavar='FILE'
+    )
+
+    parser.add_argument(
+        '--d_channels',
+        help='Instead of the values of "--channels", ' + \
+            'use these for the discriminator instead.',
+        nargs='*',
+        type=int,
+        default=[],
+        metavar='CHANNELS'
+    )
+
+    parser.add_argument(
+        '--d_skip',
+        help='Use skip connections for the discriminator. Default: %(default)s',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=False,
+        metavar='BOOL'
+    )
+
+    parser.add_argument(
+        '--d_resnet',
+        help='Use resnet connections for the discriminator. Default: %(default)s',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=True,
+        metavar='BOOL'
+    )
+
+    parser.add_argument(
+        '--d_conv_block_size',
+        help='Number of layers in a conv block in the discriminator. Default: %(default)s',
+        type=int,
+        default=2,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--d_fused_conv',
+        help='Fuse conv & downsample into a strided ' + \
+            'conv for the discriminator. Default: %(default)s',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=True,
+        metavar='BOOL'
+    )
+
+    parser.add_argument(
+        '--group_size',
+        help='Size of the groups in batch std layer. Default: %(default)s',
+        type=int,
+        default=4,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--d_activation',
+        help='The non-linear activaiton function for the discriminator. Default: %(default)s',
+        default='leaky:0.2',
+        type=str,
+        metavar='ACTIVATION'
+    )
+
+    parser.add_argument(
+        '--d_conv_resample_mode',
+        help='Resample mode for downsampling conv ' + \
+            'layers for discriminator. Default: %(default)s',
+        type=str,
+        default='FIR',
+        metavar='MODE'
+    )
+
+    parser.add_argument(
+        '--d_skip_resample_mode',
+        help='Resample mode for skip connection ' + \
+            'downsamples for the discriminator. Default: %(default)s',
+        type=str,
+        default='FIR',
+        metavar='MODE'
+    )
+
+    parser.add_argument(
+        '--d_loss',
+        help='Loss function for the disriminator. Default: %(default)s',
+        default='logistic',
+        type=str,
+        metavar='LOSS'
+    )
+
+    parser.add_argument(
+        '--d_reg',
+        help='Regularization function for the discriminator ' + \
+            'with an optional weight (:?). Default: %(default)s',
+        default='r1:10',
+        type=str,
+        metavar='REG'
+    )
+
+    parser.add_argument(
+        '--d_reg_interval',
+        help='Interval at which to regularize the discriminator. Default: %(default)s',
+        default=16,
+        type=int,
+        metavar='INTERVAL'
+    )
+
+    parser.add_argument(
+        '--d_iter',
+        help='Number of discriminator iterations per training iteration. Default: %(default)s',
+        default=1,
+        type=int,
+        metavar='ITER'
+    )
+
+    parser.add_argument(
+        '--d_lr',
+        help='The learning rate for the discriminator. Default: %(default)s',
+        default=2e-3,
+        type=float,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--d_betas',
+        help='Beta values for the discriminator Adam optimizer. Default: %(default)s',
+        type=float,
+        nargs=2,
+        default=(0, 0.99),
+        metavar='VALUE'
+    )
+
+    #----------------------------------------------------------------------------
+    # Training options
+
+    parser.add_argument(
+        '--iterations',
+        help='Number of iterations to train for. Default: %(default)s',
+        type=int,
+        default=1000000,
+        metavar='ITERATIONS'
+    )
+
+    parser.add_argument(
+        '--gpu',
+        help='The cuda device(s) to use. Example: ""--gpu 0 1" will train ' + \
+            'on GPU 0 and GPU 1. Default: Only use CPU',
+        type=int,
+        default=[],
+        nargs='*',
+        metavar='DEVICE_ID'
+    )
+
+    parser.add_argument(
+        '--distributed',
+        help='When more than one gpu device is passed, automatically ' + \
+            'start one process for each device and give it the correct ' + \
+            'distributed args (rank, world_size etc). Disable this if ' + \
+            'you want training to be performed with only one process ' + \
+            'using the DataParallel module. Default: %(default)s',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=True,
+        metavar='BOOL'
+    )
+
+    parser.add_argument(
+        '--rank',
+        help='Rank for distributed training.',
+        type=int,
+        default=None,
+    )
+
+    parser.add_argument(
+        '--world_size',
+        help='World size for distributed training.',
+        type=int,
+        default=None,
+    )
+
+    parser.add_argument(
+        '--master_addr',
+        help='Address for distributed training.',
+        type=str,
+        default=None,
+    )
+
+    parser.add_argument(
+        '--master_port',
+        help='Port for distributed training.',
+        type=str,
+        default=None,
+    )
+
+    parser.add_argument(
+        '--batch_size',
+        help='Size of each batch. Default: %(default)s',
+        default=32,
+        type=int,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--device_batch_size',
+        help='Maximum number of items to fit on single device at a time. Default: %(default)s',
+        default=4,
+        type=int,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--g_reg_batch_size',
+        help='Size of each batch used to regularize the generator. Default: %(default)s',
+        default=16,
+        type=int,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--g_reg_device_batch_size',
+        help='Maximum number of items to fit on single device when ' + \
+            'regularizing the generator. Default: %(default)s',
+        default=2,
+        type=int,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--half',
+        help='Use mixed precision training. Default: %(default)s',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=False,
+        metavar='BOOL'
+    )
+
+    parser.add_argument(
+        '--resume',
+        help='Resume from the latest saved checkpoint in the checkpoint_dir. ' + \
+            'This loads all previous training settings except for the dataset options, ' + \
+            'device args (--gpu ...) and distributed training args (--rank, --world_size e.t.c) ' + \
+            'as well as metrics and logging.',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=False,
+        metavar='BOOL'
+    )
+
+    #----------------------------------------------------------------------------
+    # Extra metric options
+
+    parser.add_argument(
+        '--fid_interval',
+        help='If specified, evaluate the FID metric with this interval.',
+        default=None,
+        type=int,
+        metavar='INTERVAL'
+    )
+
+    parser.add_argument(
+        '--ppl_interval',
+        help='If specified, evaluate the PPL metric with this interval.',
+        default=None,
+        type=int,
+        metavar='INTERVAL'
+    )
+
+    parser.add_argument(
+        '--ppl_ffhq_crop',
+        help='Crop images evaluated for PPL with crop values for FFHQ. Default: %(default)s',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=False,
+        metavar='BOOL'
+    )
+
+    #----------------------------------------------------------------------------
+    # Data options
+
+    parser.add_argument(
+        '--pixel_min',
+        help='Minimum of the value range of pixels in generated images. Default: %(default)s',
+        default=-1,
+        type=float,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--pixel_max',
+        help='Maximum of the value range of pixels in generated images. Default: %(default)s',
+        default=1,
+        type=float,
+        metavar='VALUE'
+    )
+
+    parser.add_argument(
+        '--data_channels',
+        help='Number of channels in the data. Default: 3 (RGB)',
+        default=3,
+        type=int,
+        choices=[1, 3],
+        metavar='CHANNELS'
+    )
+
+    parser.add_argument(
+        '--data_dir',
+        help='The root directory of the dataset. This argument is required!',
+        type=str,
+        default=None
+    )
+
+    parser.add_argument(
+        '--data_resize',
+        help='Resize data to fit input size of discriminator. Default: %(default)s',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=False,
+        metavar='BOOL'
+    )
+
+    parser.add_argument(
+        '--mirror_augment',
+        help='Use random horizontal flipping for data images. Default: %(default)s',
+        type=utils.bool_type,
+        const=True,
+        nargs='?',
+        default=False,
+        metavar='BOOL'
+    )
+
+    parser.add_argument(
+        '--data_workers',
+        help='Number of worker processes that handles dataloading. Default: %(default)s',
+        default=4,
+        type=int,
+        metavar='WORKERS'
+    )
+
+    #----------------------------------------------------------------------------
+    # Logging options
+
+    parser.add_argument(
+        '--checkpoint_dir',
+        help='If specified, save checkpoints to this directory.',
+        default=None,
+        type=str,
+        metavar='DIR'
+    )
+
+    parser.add_argument(
+        '--checkpoint_interval',
+        help='Save checkpoints with this interval. Default: %(default)s',
+        default=10000,
+        type=int,
+        metavar='INTERVAL'
+    )
+
+    parser.add_argument(
+        '--tensorboard_log_dir',
+        help='Log to this tensorboard directory if specified.',
+        default=None,
+        type=str,
+        metavar='DIR'
+    )
+
+    parser.add_argument(
+        '--tensorboard_image_interval',
+        help='Log images to tensorboard with this interval if specified.',
+        default=None,
+        type=int,
+        metavar='INTERVAL'
+    )
+
+    parser.add_argument(
+        '--tensorboard_image_size',
+        help='Size of images logged to tensorboard. Default: %(default)s',
+        default=256,
+        type=int,
+        metavar='VALUE'
+    )
+
+    return parser
+
+#----------------------------------------------------------------------------
+
+def get_dataset(args):
+    assert args.data_dir, '--data_dir has to be specified.'
+    height, width = [
+        shape * 2 ** (len(args.d_channels or args.channels) - 1)
+        for shape in args.base_shape
+    ]
+    dataset = utils.ImageFolder(
+        args.data_dir,
+        mirror=args.mirror_augment,
+        pixel_min=args.pixel_min,
+        pixel_max=args.pixel_max,
+        height=height,
+        width=width,
+        resize=args.data_resize,
+        grayscale=args.data_channels == 1
+    )
+    assert len(dataset), 'No images found at {}'.format(args.data_dir)
+    return dataset
+
+#----------------------------------------------------------------------------
+
+def get_models(args):
+    common_kwargs = dict(
+        data_channels=args.data_channels,
+        base_shape=args.base_shape,
+        conv_filter=args.filter,
+        skip_filter=args.filter,
+        kernel_size=args.kernel_size,
+        conv_pad_mode=args.pad_mode,
+        conv_pad_constant=args.pad_constant,
+        filter_pad_mode=args.filter_pad_mode,
+        filter_pad_constant=args.filter_pad_constant,
+        pad_once=args.pad_once,
+        weight_scale=args.weight_scale
+    )
+
+    if args.g_file:
+        G = stylegan2.models.load(args.g_file)
+        assert isinstance(G, stylegan2.models.Generator), \
+                '`--g_file` should specify a generator model, found {}'.format(type(G))
+    else:
+
+        G_M = stylegan2.models.GeneratorMapping(
+            latent_size=args.latent,
+            label_size=args.label,
+            num_layers=args.latent_mapping_layers,
+            hidden=args.latent,
+            activation=args.g_activation,
+            normalize_input=args.normalize_latent,
+            lr_mul=args.latent_mapping_lr_mul,
+            weight_scale=args.weight_scale
+        )
+
+        G_S = stylegan2.models.GeneratorSynthesis(
+            channels=args.g_channels or args.channels,
+            latent_size=args.latent,
+            demodulate=args.g_normalize,
+            modulate_data_out=args.modulate_rgb,
+            conv_block_size=args.g_conv_block_size,
+            activation=args.g_activation,
+            conv_resample_mode=args.g_conv_resample_mode,
+            skip_resample_mode=args.g_skip_resample_mode,
+            resnet=args.g_resnet,
+            skip=args.g_skip,
+            fused_resample=args.g_fused_conv,
+            **common_kwargs
+        )
+
+        G = stylegan2.models.Generator(G_mapping=G_M, G_synthesis=G_S)
+
+    if args.d_file:
+        D = stylegan2.models.load(args.d_file)
+        assert isinstance(D, stylegan2.models.Discriminator), \
+                '`--d_file` should specify a discriminator model, found {}'.format(type(D))
+    else:
+        D = stylegan2.models.Discriminator(
+            channels=args.d_channels or args.channels,
+            label_size=args.label,
+            conv_block_size=args.d_conv_block_size,
+            activation=args.d_activation,
+            conv_resample_mode=args.d_conv_resample_mode,
+            skip_resample_mode=args.d_skip_resample_mode,
+            mbstd_group_size=args.group_size,
+            resnet=args.d_resnet,
+            skip=args.d_skip,
+            fused_resample=args.d_fused_conv,
+            **common_kwargs
+        )
+    assert len(G.G_synthesis.channels) == len(D.channels), \
+        'While the number of channels for each layer can ' + \
+        'differ between generator and discriminator, the ' + \
+        'number of layers have to be the same. Received ' + \
+        '{} generator layers and {} discriminator layers.'.format(
+            len(G.G_synthesis.channels), len(D.channels))
+
+    return G, D
+
+#----------------------------------------------------------------------------
+
+def get_trainer(args):
+    dataset = get_dataset(args)
+    if args.resume and stylegan2.train._find_checkpoint(args.checkpoint_dir):
+        trainer = stylegan2.train.Trainer.load_checkpoint(
+            args.checkpoint_dir,
+            dataset,
+            device=args.gpu,
+            rank=args.rank,
+            world_size=args.world_size,
+            master_addr=args.master_addr,
+            master_port=args.master_port,
+            tensorboard_log_dir=args.tensorboard_log_dir
+        )
+    else:
+        G, D = get_models(args)
+        trainer = stylegan2.train.Trainer(
+            G=G,
+            D=D,
+            latent_size=args.latent,
+            dataset=dataset,
+            device=args.gpu,
+            batch_size=args.batch_size,
+            device_batch_size=args.device_batch_size,
+            label_size=args.label,
+            data_workers=args.data_workers,
+            G_loss=args.g_loss,
+            D_loss=args.d_loss,
+            G_reg=args.g_reg,
+            G_reg_interval=args.g_reg_interval,
+            G_opt_kwargs={'lr': args.g_lr, 'betas': args.g_betas},
+            G_reg_batch_size=args.g_reg_batch_size,
+            G_reg_device_batch_size=args.g_reg_device_batch_size,
+            D_reg=args.d_reg,
+            D_reg_interval=args.d_reg_interval,
+            D_opt_kwargs={'lr': args.d_lr, 'betas': args.d_betas},
+            style_mix_prob=args.style_mix,
+            G_iter=args.g_iter,
+            D_iter=args.d_iter,
+            tensorboard_log_dir=args.tensorboard_log_dir,
+            checkpoint_dir=args.checkpoint_dir,
+            checkpoint_interval=args.checkpoint_interval,
+            half=args.half,
+            rank=args.rank,
+            world_size=args.world_size,
+            master_addr=args.master_addr,
+            master_port=args.master_port
+        )
+    if args.fid_interval and not args.rank:
+        fid_model = inception.InceptionV3FeatureExtractor(
+            pixel_min=args.pixel_min, pixel_max=args.pixel_max)
+        trainer.register_metric(
+            name='FID (299x299)',
+            eval_fn=fid.FID(
+                trainer.Gs,
+                trainer.prior_generator,
+                dataset=dataset,
+                fid_model=fid_model,
+                fid_size=299,
+                reals_batch_size=64
+            ),
+            interval=args.fid_interval
+        )
+        trainer.register_metric(
+            name='FID',
+            eval_fn=fid.FID(
+                trainer.Gs,
+                trainer.prior_generator,
+                dataset=dataset,
+                fid_model=fid_model,
+                fid_size=None
+            ),
+            interval=args.fid_interval
+        )
+    if args.ppl_interval and not args.rank:
+        lpips_model = lpips.LPIPS_VGG16(
+            pixel_min=args.pixel_min, pixel_max=args.pixel_max)
+        crop = None
+        if args.ppl_ffhq_crop:
+            crop = ppl.PPL.FFHQ_CROP
+        trainer.register_metric(
+            name='PPL_end',
+            eval_fn=ppl.PPL(
+                trainer.Gs,
+                trainer.prior_generator,
+                full_sampling=False,
+                crop=crop,
+                lpips_model=lpips_model,
+                lpips_size=256
+            ),
+            interval=args.ppl_interval
+        )
+        trainer.register_metric(
+            name='PPL_full',
+            eval_fn=ppl.PPL(
+                trainer.Gs,
+                trainer.prior_generator,
+                full_sampling=True,
+                crop=crop,
+                lpips_model=lpips_model,
+                lpips_size=256
+            ),
+            interval=args.ppl_interval
+        )
+    if args.tensorboard_image_interval:
+        for static in [True, False]:
+            for trunc in [0.5, 0.7, 1.0]:
+                if static:
+                    name = 'static'
+                else:
+                    name = 'random'
+                name += '/trunc_{:.1f}'.format(trunc)
+                trainer.add_tensorboard_image_logging(
+                    name=name,
+                    num_images=4,
+                    interval=args.tensorboard_image_interval,
+                    resize=args.tensorboard_image_size,
+                    seed=1234567890 if static else None,
+                    truncation_psi=trunc,
+                    pixel_min=args.pixel_min,
+                    pixel_max=args.pixel_max
+                )
+    return trainer
+
+#----------------------------------------------------------------------------
+
+def run(args):
+    if not args.rank:
+        if not (args.checkpoint_dir or args.output):
+            warnings.warn(
+                'Neither an output path or checkpoint dir has been ' + \
+                'given. Weights from this training run will never ' + \
+                'be saved.'
+            )
+        if args.output:
+            assert os.path.isdir(args.output) or not os.path.splitext(args.output)[-1], \
+                '--output argument should specify a directory, not a file.'
+    trainer = get_trainer(args)
+    trainer.train(iterations=args.iterations)
+    if not args.rank and args.output:
+        print('Saving models to {}'.format(args.output))
+        if not os.path.exists(args.output):
+            os.makedirs(args.output)
+        for model_name in ['G', 'D', 'Gs']:
+            getattr(trainer, model_name).save(
+                os.path.join(args.output_dir, model_name + '.pth'))
+
+#----------------------------------------------------------------------------
+
+def run_distributed(rank, args):
+    args.rank = rank
+    args.world_size = len(args.gpu)
+    args.gpu = args.gpu[rank]
+    args.master_addr = args.master_addr or '127.0.0.1'
+    args.master_port = args.master_port or '23456'
+    run(args)
+
+#----------------------------------------------------------------------------
+
+def main():
+    parser = get_arg_parser()
+    args = parser.parse_args()
+    if len(args.gpu) > 1 and args.distributed:
+        assert args.rank is None and args.world_size is None, \
+            'When --distributed is enabled (default) the rank and ' + \
+            'world size can not be given as this is set up automatically. ' + \
+            'Use --distributed 0 to disable automatic setup of distributed training.'
+        mp.spawn(run_distributed, nprocs=len(args.gpu), args=(args,))
+    else:
+        run(args)
+
+#----------------------------------------------------------------------------
+
+if __name__ == '__main__':
+    main()

stylegan2/__init__.py

@@ -0,0 +1,5 @@
+from . import external_models
+from . import metrics
+from . import models
+from . import project
+from . import train

stylegan2/external_models/__init__.py

@@ -0,0 +1,2 @@
+from . import inception
+from . import lpips

stylegan2/external_models/inception.py

@@ -0,0 +1,276 @@
+"""
+Code adapted from https://github.com/mseitzer/pytorch-fid/
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+   http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import models
+
+try:
+    from torchvision.models.utils import load_state_dict_from_url
+except ImportError:
+    from torch.utils.model_zoo import load_url as load_state_dict_from_url
+
+# Inception weights ported to Pytorch from
+# http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
+FID_WEIGHTS_URL = 'https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth'
+
+
+class InceptionV3FeatureExtractor(nn.Module):
+    """Pretrained InceptionV3 network returning feature maps"""
+
+    # Index of default block of inception to return,
+    # corresponds to output of final average pooling
+    DEFAULT_BLOCK_INDEX = 3
+
+    # Maps feature dimensionality to their output blocks indices
+    BLOCK_INDEX_BY_DIM = {
+        64: 0,   # First max pooling features
+        192: 1,  # Second max pooling featurs
+        768: 2,  # Pre-aux classifier features
+        2048: 3  # Final average pooling features
+    }
+
+    def __init__(self,
+                 output_block=DEFAULT_BLOCK_INDEX,
+                 pixel_min=-1,
+                 pixel_max=1):
+        """
+        Build pretrained InceptionV3
+        Arguments:
+            output_block (int): Index of block to return features of.
+                Possible values are:
+                    - 0: corresponds to output of first max pooling
+                    - 1: corresponds to output of second max pooling
+                    - 2: corresponds to output which is fed to aux classifier
+                    - 3: corresponds to output of final average pooling
+            pixel_min (float): Min value for inputs. Default value is -1.
+            pixel_max (float): Max value for inputs. Default value is 1.
+        """
+        super(InceptionV3FeatureExtractor, self).__init__()
+
+        assert 0 <= output_block <= 3, '`output_block` can only be ' + \
+            '0 <= `output_block` <= 3.'
+
+        inception = fid_inception_v3()
+
+        blocks = []
+
+        # Block 0: input to maxpool1
+        block0 = [
+            inception.Conv2d_1a_3x3,
+            inception.Conv2d_2a_3x3,
+            inception.Conv2d_2b_3x3,
+            nn.MaxPool2d(kernel_size=3, stride=2)
+        ]
+        blocks.append(nn.Sequential(*block0))
+
+        # Block 1: maxpool1 to maxpool2
+        if output_block >= 1:
+            block1 = [
+                inception.Conv2d_3b_1x1,
+                inception.Conv2d_4a_3x3,
+                nn.MaxPool2d(kernel_size=3, stride=2)
+            ]
+            blocks.append(nn.Sequential(*block1))
+
+        # Block 2: maxpool2 to aux classifier
+        if output_block >= 2:
+            block2 = [
+                inception.Mixed_5b,
+                inception.Mixed_5c,
+                inception.Mixed_5d,
+                inception.Mixed_6a,
+                inception.Mixed_6b,
+                inception.Mixed_6c,
+                inception.Mixed_6d,
+                inception.Mixed_6e,
+            ]
+            blocks.append(nn.Sequential(*block2))
+
+        # Block 3: aux classifier to final avgpool
+        if output_block >= 3:
+            block3 = [
+                inception.Mixed_7a,
+                inception.Mixed_7b,
+                inception.Mixed_7c,
+                nn.AdaptiveAvgPool2d(output_size=(1, 1))
+            ]
+            blocks.append(nn.Sequential(*block3))
+
+        self.main = nn.Sequential(*blocks)
+        self.pixel_nin = pixel_min
+        self.pixel_max = pixel_max
+        self.requires_grad_(False)
+        self.eval()
+
+    def _scale(self, x):
+        if self.pixel_min != -1 or self.pixel_max != 1:
+            x = (2*x - self.pixel_min - self.pixel_max) \
+                / (self.pixel_max - self.pixel_min)
+        return x
+
+    def forward(self, input):
+        """
+        Get Inception feature maps.
+        Arguments:
+            input (torch.Tensor)
+        Returns:
+            feature_maps (torch.Tensor)
+        """
+        return self.main(input)
+
+
+def fid_inception_v3():
+    """Build pretrained Inception model for FID computation
+    The Inception model for FID computation uses a different set of weights
+    and has a slightly different structure than torchvision's Inception.
+    This method first constructs torchvision's Inception and then patches the
+    necessary parts that are different in the FID Inception model.
+    """
+    inception = models.inception_v3(num_classes=1008,
+                                    aux_logits=False,
+                                    pretrained=False)
+    inception.Mixed_5b = FIDInceptionA(192, pool_features=32)
+    inception.Mixed_5c = FIDInceptionA(256, pool_features=64)
+    inception.Mixed_5d = FIDInceptionA(288, pool_features=64)
+    inception.Mixed_6b = FIDInceptionC(768, channels_7x7=128)
+    inception.Mixed_6c = FIDInceptionC(768, channels_7x7=160)
+    inception.Mixed_6d = FIDInceptionC(768, channels_7x7=160)
+    inception.Mixed_6e = FIDInceptionC(768, channels_7x7=192)
+    inception.Mixed_7b = FIDInceptionE_1(1280)
+    inception.Mixed_7c = FIDInceptionE_2(2048)
+
+    state_dict = load_state_dict_from_url(FID_WEIGHTS_URL, progress=True)
+    inception.load_state_dict(state_dict)
+    return inception
+
+
+class FIDInceptionA(models.inception.InceptionA):
+    """InceptionA block patched for FID computation"""
+    def __init__(self, in_channels, pool_features):
+        super(FIDInceptionA, self).__init__(in_channels, pool_features)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch5x5 = self.branch5x5_1(x)
+        branch5x5 = self.branch5x5_2(branch5x5)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)
+
+        # Patch: Tensorflow's average pool does not use the padded zero's in
+        # its average calculation
+        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                   count_include_pad=False)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool]
+        return torch.cat(outputs, 1)
+
+
+class FIDInceptionC(models.inception.InceptionC):
+    """InceptionC block patched for FID computation"""
+    def __init__(self, in_channels, channels_7x7):
+        super(FIDInceptionC, self).__init__(in_channels, channels_7x7)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch7x7 = self.branch7x7_1(x)
+        branch7x7 = self.branch7x7_2(branch7x7)
+        branch7x7 = self.branch7x7_3(branch7x7)
+
+        branch7x7dbl = self.branch7x7dbl_1(x)
+        branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl)
+
+        # Patch: Tensorflow's average pool does not use the padded zero's in
+        # its average calculation
+        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                   count_include_pad=False)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool]
+        return torch.cat(outputs, 1)
+
+
+class FIDInceptionE_1(models.inception.InceptionE):
+    """First InceptionE block patched for FID computation"""
+    def __init__(self, in_channels):
+        super(FIDInceptionE_1, self).__init__(in_channels)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch3x3 = self.branch3x3_1(x)
+        branch3x3 = [
+            self.branch3x3_2a(branch3x3),
+            self.branch3x3_2b(branch3x3),
+        ]
+        branch3x3 = torch.cat(branch3x3, 1)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = [
+            self.branch3x3dbl_3a(branch3x3dbl),
+            self.branch3x3dbl_3b(branch3x3dbl),
+        ]
+        branch3x3dbl = torch.cat(branch3x3dbl, 1)
+
+        # Patch: Tensorflow's average pool does not use the padded zero's in
+        # its average calculation
+        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                   count_include_pad=False)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
+        return torch.cat(outputs, 1)
+
+
+class FIDInceptionE_2(models.inception.InceptionE):
+    """Second InceptionE block patched for FID computation"""
+    def __init__(self, in_channels):
+        super(FIDInceptionE_2, self).__init__(in_channels)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch3x3 = self.branch3x3_1(x)
+        branch3x3 = [
+            self.branch3x3_2a(branch3x3),
+            self.branch3x3_2b(branch3x3),
+        ]
+        branch3x3 = torch.cat(branch3x3, 1)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = [
+            self.branch3x3dbl_3a(branch3x3dbl),
+            self.branch3x3dbl_3b(branch3x3dbl),
+        ]
+        branch3x3dbl = torch.cat(branch3x3dbl, 1)
+
+        # Patch: The FID Inception model uses max pooling instead of average
+        # pooling. This is likely an error in this specific Inception
+        # implementation, as other Inception models use average pooling here
+        # (which matches the description in the paper).
+        branch_pool = F.max_pool2d(x, kernel_size=3, stride=1, padding=1)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
+        return torch.cat(outputs, 1)

stylegan2/external_models/lpips.py

@@ -0,0 +1,78 @@
+"""
+Code adapted from https://github.com/richzhang/PerceptualSimilarity
+
+Original License:
+Copyright (c) 2018, Richard Zhang, Phillip Isola, Alexei A. Efros, Eli Shechtman, Oliver Wang
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+"""
+import torch
+from torch import nn
+import torchvision
+
+
+class LPIPS_VGG16(nn.Module):
+    _FEATURE_IDX = [0, 4, 9, 16, 23, 30]
+    _LINEAR_WEIGHTS_URL = 'https://github.com/richzhang/PerceptualSimilarity' + \
+                          '/blob/master/lpips/weights/v0.1/vgg.pth?raw=true'
+
+    def __init__(self, pixel_min=-1, pixel_max=1):
+        super(LPIPS_VGG16, self).__init__()
+        features = torchvision.models.vgg16(pretrained=True).features
+        self.slices = nn.ModuleList()
+        linear_weights = torch.utils.model_zoo.load_url(self._LINEAR_WEIGHTS_URL)
+        for i in range(1, len(self._FEATURE_IDX)):
+            idx_range = range(self._FEATURE_IDX[i - 1], self._FEATURE_IDX[i])
+            self.slices.append(nn.Sequential(*[features[j] for j in idx_range]))
+        self.linear_layers = nn.ModuleList()
+        for weight in torch.utils.model_zoo.load_url(self._LINEAR_WEIGHTS_URL).values():
+            weight = weight.view(1, -1)
+            linear = nn.Linear(weight.size(1), 1, bias=False)
+            linear.weight.data.copy_(weight)
+            self.linear_layers.append(linear)
+        self.register_buffer('shift', torch.Tensor([-.030,-.088,-.188]).view(1, -1, 1, 1))
+        self.register_buffer('scale', torch.Tensor([.458,.448,.450]).view(1, -1, 1, 1))
+        self.pixel_min = pixel_min
+        self.pixel_max = pixel_max
+        self.requires_grad_(False)
+        self.eval()
+
+    def _scale(self, x):
+        if self.pixel_min != -1 or self.pixel_max != 1:
+            x = (2*x - self.pixel_min - self.pixel_max) \
+                / (self.pixel_max - self.pixel_min)
+        return (x - self.shift) / self.scale
+
+    @staticmethod
+    def _normalize_tensor(feature_maps, eps=1e-8):
+        rnorm = torch.rsqrt(torch.sum(feature_maps ** 2, dim=1, keepdim=True) + eps)
+        return feature_maps * rnorm
+
+    def forward(self, x0, x1, eps=1e-8):
+        x0, x1 = self._scale(x0), self._scale(x1)
+        dist = 0
+        for slice, linear in zip(self.slices, self.linear_layers):
+            x0, x1 = slice(x0), slice(x1)
+            _x0, _x1 = self._normalize_tensor(x0, eps), self._normalize_tensor(x1, eps)
+            dist += linear(torch.mean((_x0 - _x1) ** 2, dim=[-1, -2]))
+        return dist.view(-1)

stylegan2/loss_fns.py

@@ -0,0 +1,347 @@
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+from . import utils
+
+
+def _grad(input, output, retain_graph):
+    # https://discuss.pytorch.org/t/gradient-penalty-loss-with-modified-weights/64910
+    # Currently not possible to not
+    # retain graph for regularization losses.
+    # Ugly hack is to always set it to True.
+    retain_graph = True
+    grads = torch.autograd.grad(
+        output.sum(),
+        input,
+        only_inputs=True,
+        retain_graph=retain_graph,
+        create_graph=True
+    )
+    return grads[0]
+
+
+def _grad_pen(input, output, gamma, constraint=1, onesided=False, retain_graph=True):
+    grad = _grad(input, output, retain_graph=retain_graph)
+    grad = grad.view(grad.size(0), -1)
+    grad_norm = grad.norm(2, dim=1)
+    if onesided:
+        gp = torch.max(0, grad_norm - constraint)
+    else:
+        gp = (grad_norm - constraint) ** 2
+    return gamma * gp.mean()
+
+
+def _grad_reg(input, output, gamma, retain_graph=True):
+    grad = _grad(input, output, retain_graph=retain_graph)
+    grad = grad.view(grad.size(0), -1)
+    gr = (grad ** 2).sum(1)
+    return (0.5 * gamma) * gr.mean()
+
+
+def _pathreg(dlatents, fakes, pl_avg, pl_decay, gamma, retain_graph=True):
+    retain_graph = True
+    pl_noise = torch.empty_like(fakes).normal_().div_(np.sqrt(np.prod(fakes.size()[2:])))
+    pl_grad = _grad(dlatents, torch.sum(pl_noise * fakes), retain_graph=retain_graph)
+    pl_length = torch.sqrt(torch.mean(torch.sum(pl_grad ** 2, dim=2), dim=1))
+    with torch.no_grad():
+        pl_avg.add_(pl_decay * (torch.mean(pl_length) - pl_avg))
+    return gamma * torch.mean((pl_length - pl_avg) ** 2)
+
+
+#----------------------------------------------------------------------------
+# Logistic loss from the paper
+# "Generative Adversarial Nets", Goodfellow et al. 2014
+
+
+def G_logistic(G,
+               D,
+               latents,
+               latent_labels=None,
+               *args,
+               **kwargs):
+    fake_scores = D(G(latents, labels=latent_labels), labels=latent_labels).float()
+    loss = - F.binary_cross_entropy_with_logits(fake_scores, torch.zeros_like(fake_scores))
+    reg = None
+    return loss, reg
+
+
+def G_logistic_ns(G,
+                  D,
+                  latents,
+                  latent_labels=None,
+                  *args,
+                  **kwargs):
+    fake_scores = D(G(latents, labels=latent_labels), labels=latent_labels).float()
+    loss = F.binary_cross_entropy_with_logits(fake_scores, torch.ones_like(fake_scores))
+    reg = None
+    return loss, reg
+
+
+def D_logistic(G,
+               D,
+               latents,
+               reals,
+               latent_labels=None,
+               real_labels=None,
+               *args,
+               **kwargs):
+    assert (latent_labels is None) == (real_labels is None)
+    with torch.no_grad():
+        fakes = G(latents, labels=latent_labels)
+    real_scores = D(reals, labels=real_labels).float()
+    fake_scores = D(fakes, labels=latent_labels).float()
+    real_loss = F.binary_cross_entropy_with_logits(real_scores, torch.ones_like(real_scores))
+    fake_loss = F.binary_cross_entropy_with_logits(fake_scores, torch.zeros_like(fake_scores))
+    loss = real_loss + fake_loss
+    reg = None
+    return loss, reg
+
+
+#----------------------------------------------------------------------------
+# R1 and R2 regularizers from the paper
+# "Which Training Methods for GANs do actually Converge?", Mescheder et al. 2018
+
+
+def D_r1(D,
+         reals,
+         real_labels=None,
+         gamma=10,
+         *args,
+         **kwargs):
+    loss = None
+    reg = None
+    if gamma:
+        reals.requires_grad_(True)
+        real_scores = D(reals, labels=real_labels)
+        reg = _grad_reg(
+            input=reals, output=real_scores, gamma=gamma, retain_graph=False).float()
+    return loss, reg
+
+
+def D_r2(D,
+         G,
+         latents,
+         latent_labels=None,
+         gamma=10,
+         *args,
+         **kwargs):
+    loss = None
+    reg = None
+    if gamma:
+        with torch.no_grad():
+            fakes = G(latents, labels=latent_labels)
+        fakes.requires_grad_(True)
+        fake_scores = D(fakes, labels=latent_labels)
+        reg = _grad_reg(
+            input=fakes, output=fake_scores, gamma=gamma, retain_graph=False).float()
+    return loss, reg
+
+
+def D_logistic_r1(G,
+                  D,
+                  latents,
+                  reals,
+                  latent_labels=None,
+                  real_labels=None,
+                  gamma=10,
+                  *args,
+                  **kwargs):
+    assert (latent_labels is None) == (real_labels is None)
+    with torch.no_grad():
+        fakes = G(latents, labels=latent_labels)
+    if gamma:
+        reals.requires_grad_(True)
+    real_scores = D(reals, labels=real_labels).float()
+    fake_scores = D(fakes, labels=latent_labels).float()
+    real_loss = F.binary_cross_entropy_with_logits(real_scores, torch.ones_like(real_scores))
+    fake_loss = F.binary_cross_entropy_with_logits(fake_scores, torch.zeros_like(fake_scores))
+    loss = real_loss + fake_loss
+    reg = None
+    if gamma:
+        reg = _grad_reg(
+            input=reals, output=real_scores, gamma=gamma, retain_graph=True).float()
+    return loss, reg
+
+
+def D_logistic_r2(G,
+                  D,
+                  latents,
+                  reals,
+                  latent_labels=None,
+                  real_labels=None,
+                  gamma=10,
+                  *args,
+                  **kwargs):
+    assert (latent_labels is None) == (real_labels is None)
+    with torch.no_grad():
+        fakes = G(latents, labels=latent_labels)
+    if gamma:
+        fakes.requires_grad_(True)
+    real_scores = D(reals, labels=real_labels).float()
+    fake_scores = D(fakes, labels=latent_labels).float()
+    real_loss = F.binary_cross_entropy_with_logits(real_scores, torch.ones_like(real_scores))
+    fake_loss = F.binary_cross_entropy_with_logits(fake_scores, torch.zeros_like(fake_scores))
+    loss = real_loss + fake_loss
+    reg = None
+    if gamma:
+        reg = _grad_reg(
+            input=fakes, output=fake_scores, gamma=gamma, retain_graph=True).float()
+    return loss, reg
+
+
+#----------------------------------------------------------------------------
+# Non-saturating logistic loss with path length regularizer from the paper
+# "Analyzing and Improving the Image Quality of StyleGAN", Karras et al. 2019
+
+
+def G_pathreg(G,
+              latents,
+              pl_avg,
+              latent_labels=None,
+              pl_decay=0.01,
+              gamma=2,
+              *args,
+              **kwargs):
+    loss = None
+    reg = None
+    if gamma:
+        fakes, dlatents = G(latents, labels=latent_labels, return_dlatents=True, mapping_grad=False)
+        reg = _pathreg(
+            dlatents=dlatents,
+            fakes=fakes,
+            pl_avg=pl_avg,
+            pl_decay=pl_decay,
+            gamma=gamma,
+            retain_graph=False
+        ).float()
+    return loss, reg
+
+
+def G_logistic_ns_pathreg(G,
+                          D,
+                          latents,
+                          pl_avg,
+                          latent_labels=None,
+                          pl_decay=0.01,
+                          gamma=2,
+                          *args,
+                          **kwargs):
+    fakes, dlatents = G(latents, labels=latent_labels, return_dlatents=True)
+    fake_scores = D(fakes, labels=latent_labels).float()
+    loss = F.binary_cross_entropy_with_logits(fake_scores, torch.ones_like(fake_scores))
+    reg = None
+    if gamma:
+        reg = _pathreg(
+            dlatents=dlatents,
+            fakes=fakes,
+            pl_avg=pl_avg,
+            pl_decay=pl_decay,
+            gamma=gamma,
+            retain_graph=True
+        ).float()
+    return loss, reg
+
+
+#----------------------------------------------------------------------------
+# WGAN loss from the paper
+# "Wasserstein Generative Adversarial Networks", Arjovsky et al. 2017
+
+
+def G_wgan(G,
+           D,
+           latents,
+           latent_labels=None,
+           *args,
+           **kwargs):
+    fake_scores = D(G(latents, labels=latent_labels), labels=latent_labels).float()
+    loss = -fake_scores.mean()
+    reg = None
+    return loss, reg
+
+
+def D_wgan(G,
+           D,
+           latents,
+           reals,
+           latent_labels=None,
+           real_labels=None,
+           drift_gamma=0.001,
+           *args,
+           **kwargs):
+    assert (latent_labels is None) == (real_labels is None)
+    with torch.no_grad():
+        fakes = G(latents, labels=latent_labels)
+    real_scores = D(reals, labels=real_labels).float()
+    fake_scores = D(fakes, labels=latent_labels).float()
+    loss = fake_scores.mean() - real_scores.mean()
+    if drift_gamma:
+        loss += drift_gamma * torch.mean(real_scores ** 2)
+    reg = None
+    return loss, reg
+
+
+#----------------------------------------------------------------------------
+# WGAN-GP loss from the paper
+# "Improved Training of Wasserstein GANs", Gulrajani et al. 2017
+
+
+def D_gp(G,
+         D,
+         latents,
+         reals,
+         latent_labels=None,
+         real_labels=None,
+         gamma=0,
+         constraint=1,
+         *args,
+         **kwargs):
+    loss = None
+    reg = None
+    if gamma:
+        assert (latent_labels is None) == (real_labels is None)
+        with torch.no_grad():
+            fakes = G(latents, labels=latent_labels)
+        assert reals.size() == fakes.size()
+        if latent_labels:
+            assert latent_labels == real_labels
+        alpha = torch.empty(reals.size(0)).uniform_()
+        alpha = alpha.view(-1, *[1] * (reals.dim() - 1))
+        interp = utils.lerp(reals, fakes, alpha).requires_grad_(True)
+        interp_scores = D(interp, labels=latent_labels)
+        reg = _grad_pen(
+            input=interp, output=interp_scores, gamma=gamma, retain_graph=False).float()
+    return loss, reg
+
+
+def D_wgan_gp(G,
+              D,
+              latents,
+              reals,
+              latent_labels=None,
+              real_labels=None,
+              gamma=0,
+              drift_gamma=0.001,
+              constraint=1,
+              *args,
+              **kwargs):
+    assert (latent_labels is None) == (real_labels is None)
+    with torch.no_grad():
+        fakes = G(latents, labels=latent_labels)
+    real_scores = D(reals, labels=real_labels).float()
+    fake_scores = D(fakes, labels=latent_labels).float()
+    loss = fake_scores.mean() - real_scores.mean()
+    if drift_gamma:
+        loss += drift_gamma * torch.mean(real_scores ** 2)
+    reg = None
+    if gamma:
+        assert reals.size() == fakes.size()
+        if latent_labels:
+            assert latent_labels == real_labels
+        alpha = torch.empty(reals.size(0)).uniform_()
+        alpha = alpha.view(-1, *[1] * (reals.dim() - 1))
+        interp = utils.lerp(reals, fakes, alpha).requires_grad_(True)
+        interp_scores = D(interp, labels=latent_labels)
+        reg = _grad_pen(
+            input=interp, output=interp_scores, gamma=gamma, retain_graph=True).float()
+    return loss, reg

stylegan2/metrics/__init__.py

@@ -0,0 +1,2 @@
+from . import fid
+from . import ppl

stylegan2/metrics/fid.py

@@ -0,0 +1,210 @@
+import warnings
+import numbers
+import numpy as np
+import scipy
+import torch
+from torch.nn import functional as F
+
+from .. import models, utils
+from ..external_models import inception
+
+
+class _TruncatedDataset:
+    """
+    Truncates a dataset, making only part of it accessible
+    by `torch.utils.data.DataLoader`.
+    """
+
+    def __init__(self, dataset, max_len):
+        self.dataset = dataset
+        self.max_len = max_len
+
+    def __len__(self):
+        return min(len(self.dataset), self.max_len)
+
+    def __getitem__(self, index):
+        return self.dataset[index]
+
+
+class FID:
+    """
+    This class evaluates the FID metric of a generator.
+    Arguments:
+        G (Generator)
+        prior_generator (PriorGenerator)
+        dataset (indexable)
+        device (int, str, torch.device, optional): The device
+            to use for calculations. By default, the same device
+            is chosen as the parameters in `generator` reside on.
+        num_samples (int): Number of samples of reals and fakes
+            to gather statistics for which are used for calculating
+            the metric. Default value is 50 000.
+        fid_model (nn.Module): A model that returns feature maps
+            of shape (batch_size, features, *). Default value
+            is InceptionV3.
+        fid_size (int, optional): Resize any data fed to `fid_model` by scaling
+            the data so that its smallest side is the same size as this
+            argument.
+        truncation_psi (float, optional): Truncation of the generator
+            when evaluating.
+        truncation_cutoff (int, optional): Cutoff for truncation when
+            evaluating.
+        reals_batch_size (int, optional): Batch size to use for real
+            samples statistics gathering.
+        reals_data_workers (int, optional): Number of workers fetching
+            the real data samples. Default value is 0.
+        verbose (bool): Write progress of gathering statistics for reals
+            to stdout. Default value is True.
+    """
+    def __init__(self,
+                 G,
+                 prior_generator,
+                 dataset,
+                 device=None,
+                 num_samples=50000,
+                 fid_model=None,
+                 fid_size=None,
+                 truncation_psi=None,
+                 truncation_cutoff=None,
+                 reals_batch_size=None,
+                 reals_data_workers=0,
+                 verbose=True):
+        device_ids = []
+        if isinstance(G, torch.nn.DataParallel):
+            device_ids = G.device_ids
+        G = utils.unwrap_module(G)
+        assert isinstance(G, models.Generator)
+        assert isinstance(prior_generator, utils.PriorGenerator)
+        if device is None:
+            device = next(G.parameters()).device
+        else:
+            device = torch.device(device)
+        assert torch.device(prior_generator.device) == device, \
+            'Prior generator device ({}) '.format(torch.device(prior_generator)) + \
+            'is not the same as the specified (or infered from the model)' + \
+            'device ({}) for the PPL evaluation.'.format(device)
+        G.eval().to(device)
+        if device_ids:
+            G = torch.nn.DataParallel(G, device_ids=device_ids)
+        self.G = G
+        self.prior_generator = prior_generator
+        self.device = device
+        self.num_samples = num_samples
+        self.batch_size = self.prior_generator.batch_size
+        if fid_model is None:
+            warnings.warn(
+                'Using default fid model metric based on Inception V3. ' + \
+                'This metric will only work on image data where values are in ' + \
+                'the range [-1, 1], please specify another module if you want ' + \
+                'to use other kinds of data formats.'
+            )
+            fid_model = inception.InceptionV3FeatureExtractor(pixel_min=-1, pixel_max=1)
+            if device_ids:
+                fid_model = torch.nn.DataParallel(fid_model, device_ids)
+        self.fid_model = fid_model.eval().to(device)
+        self.fid_size = fid_size
+
+        dataset = _TruncatedDataset(dataset, self.num_samples)
+        dataloader = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=reals_batch_size or self.batch_size,
+            num_workers=reals_data_workers
+        )
+        features = []
+        self.labels = []
+
+        if verbose:
+            progress = utils.ProgressWriter(
+                np.ceil(self.num_samples / (reals_batch_size or self.batch_size)))
+            progress.write('FID: Gathering statistics for reals...', step=False)
+
+        for batch in dataloader:
+            data = batch
+            if isinstance(batch, (tuple, list)):
+                data = batch[0]
+                if len(batch) > 1:
+                    self.labels.append(batch[1])
+            data = self._scale_for_fid(data).to(self.device)
+            with torch.no_grad():
+                batch_features = self.fid_model(data)
+            batch_features = batch_features.view(*batch_features.size()[:2], -1).mean(-1)
+            features.append(batch_features.cpu())
+            progress.step()
+
+        if verbose:
+            progress.write('FID: Statistics for reals gathered!', step=False)
+            progress.close()
+
+        features = torch.cat(features, dim=0).numpy()
+
+        self.mu_real = np.mean(features, axis=0)
+        self.sigma_real = np.cov(features, rowvar=False)
+        self.truncation_psi = truncation_psi
+        self.truncation_cutoff = truncation_cutoff
+
+    def _scale_for_fid(self, data):
+        if not self.fid_size:
+            return data
+        scale_factor = self.fid_size / min(data.size()[2:])
+        if scale_factor == 1:
+            return data
+        mode = 'nearest'
+        if scale_factor < 1:
+            mode = 'area'
+        return F.interpolate(data, scale_factor=scale_factor, mode=mode)
+
+    def __call__(self, *args, **kwargs):
+        return self.evaluate(*args, **kwargs)
+
+    def evaluate(self, verbose=True):
+        """
+        Evaluate the FID.
+        Arguments:
+            verbose (bool): Write progress to stdout.
+                Default value is True.
+        Returns:
+            fid (float): Metric value.
+        """
+        utils.unwrap_module(self.G).set_truncation(
+            truncation_psi=self.truncation_psi, truncation_cutoff=self.truncation_cutoff)
+        self.G.eval()
+        features = []
+
+        if verbose:
+            progress = utils.ProgressWriter(np.ceil(self.num_samples / self.batch_size))
+            progress.write('FID: Gathering statistics for fakes...', step=False)
+
+        remaining = self.num_samples
+        for i in range(0, self.num_samples, self.batch_size):
+
+            latents, latent_labels = self.prior_generator(
+                batch_size=min(self.batch_size, remaining))
+            if latent_labels is not None and self.labels:
+                latent_labels = self.labels[i].to(self.device)
+                length = min(len(latents), len(latent_labels))
+                latents, latent_labels = latents[:length], latent_labels[:length]
+
+            with torch.no_grad():
+                fakes = self.G(latents, labels=latent_labels)
+
+            with torch.no_grad():
+                batch_features = self.fid_model(fakes)
+            batch_features = batch_features.view(*batch_features.size()[:2], -1).mean(-1)
+            features.append(batch_features.cpu())
+
+            remaining -= len(latents)
+            progress.step()
+
+        if verbose:
+            progress.write('FID: Statistics for fakes gathered!', step=False)
+            progress.close()
+
+        features = torch.cat(features, dim=0).numpy()
+
+        mu_fake = np.mean(features, axis=0)
+        sigma_fake = np.cov(features, rowvar=False)
+
+        m = np.square(mu_fake - self.mu_real).sum()
+        s, _ = scipy.linalg.sqrtm(np.dot(sigma_fake, self.sigma_real), disp=False)
+        dist = m + np.trace(sigma_fake + self.sigma_real - 2*s)
+        return float(np.real(dist))

stylegan2/metrics/ppl.py

@@ -0,0 +1,229 @@
+import warnings
+import numbers
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+from .. import models, utils
+from ..external_models import lpips
+
+
+class PPL:
+    """
+    This class evaluates the PPL metric of a generator.
+    Arguments:
+        G (Generator)
+        prior_generator (PriorGenerator)
+        device (int, str, torch.device, optional): The device
+            to use for calculations. By default, the same device
+            is chosen as the parameters in `generator` reside on.
+        num_samples (int): Number of samples of reals and fakes
+            to gather statistics for which are used for calculating
+            the metric. Default value is 50 000.
+        epsilon (float): Perturbation value. Default value is 1e-4.
+        use_dlatent (bool): Measure PPL against the dlatents instead
+            of the latents. Default value is True.
+        full_sampling (bool): Measure on a random interpolation between
+            two inputs. Default value is False.
+        crop (float, list, optional): Crop values that should be in the
+            range [0, 1] with 1 representing the entire data length.
+            If single value this will be the amount cropped from all
+            sides of the data. If a list of same length as number of
+            data dimensions, each crop is mirrored to both sides of
+            each respective dimension. If the length is 2 * number
+            of dimensions the crop values for the start and end of
+            a dimension may be different.
+            Example 1:
+                We have 1d data of length 10. We want to crop 1
+                from the start and end of the data. We then need
+                to use `crop=0.1` or `crop=[0.1]` or `crop=[0.1, 0.9]`.
+            Example 2:
+                We have 2d data (images) of size 10, 10 (height, width)
+                and we want to use only the top left quarter of the image
+                we would use `crop=[0, 0.5, 0, 0.5]`.
+        lpips_model (nn.Module): A model that returns feature the distance
+            between two inputs. Default value is the LPIPS VGG16 model.
+        lpips_size (int, optional): Resize any data fed to `lpips_model` by scaling
+            the data so that its smallest side is the same size as this
+            argument. Only has a default value of 256 if `lpips_model` is unspecified.
+    """
+    FFHQ_CROP = [1/8 * 3, 1/8 * 7, 1/8 * 2, 1/8 * 6]
+
+    def __init__(self,
+                 G,
+                 prior_generator,
+                 device=None,
+                 num_samples=50000,
+                 epsilon=1e-4,
+                 use_dlatent=True,
+                 full_sampling=False,
+                 crop=None,
+                 lpips_model=None,
+                 lpips_size=None):
+        device_ids = []
+        if isinstance(G, torch.nn.DataParallel):
+            device_ids = G.device_ids
+        G = utils.unwrap_module(G)
+        assert isinstance(G, models.Generator)
+        assert isinstance(prior_generator, utils.PriorGenerator)
+        if device is None:
+            device = next(G.parameters()).device
+        else:
+            device = torch.device(device)
+        assert torch.device(prior_generator.device) == device, \
+            'Prior generator device ({}) '.format(torch.device(prior_generator)) + \
+            'is not the same as the specified (or infered from the model)' + \
+            'device ({}) for the PPL evaluation.'.format(device)
+        G.eval().to(device)
+        self.G_mapping = G.G_mapping
+        self.G_synthesis = G.G_synthesis
+        if device_ids:
+            self.G_mapping = torch.nn.DataParallel(self.G_mapping, device_ids=device_ids)
+            self.G_synthesis = torch.nn.DataParallel(self.G_synthesis, device_ids=device_ids)
+        self.prior_generator = prior_generator
+        self.device = device
+        self.num_samples = num_samples
+        self.epsilon = epsilon
+        self.use_dlatent = use_dlatent
+        self.full_sampling = full_sampling
+        self.crop = crop
+        self.batch_size = self.prior_generator.batch_size
+        if lpips_model is None:
+            warnings.warn(
+                'Using default LPIPS distance metric based on VGG 16. ' + \
+                'This metric will only work on image data where values are in ' + \
+                'the range [-1, 1], please specify an lpips module if you want ' + \
+                'to use other kinds of data formats.'
+            )
+            lpips_model = lpips.LPIPS_VGG16(pixel_min=-1, pixel_max=1)
+            if device_ids:
+                lpips_model = torch.nn.DataParallel(lpips_model, device_ids=device_ids)
+            lpips_size = lpips_size or 256
+        self.lpips_model = lpips_model.eval().to(device)
+        self.lpips_size = lpips_size
+
+    def _scale_for_lpips(self, data):
+        if not self.lpips_size:
+            return data
+        scale_factor = self.lpips_size / min(data.size()[2:])
+        if scale_factor == 1:
+            return data
+        mode = 'nearest'
+        if scale_factor < 1:
+            mode = 'area'
+        return F.interpolate(data, scale_factor=scale_factor, mode=mode)
+
+    def crop_data(self, data):
+        if not self.crop:
+            return data
+        dim = data.dim() - 2
+        if isinstance(self.crop, numbers.Number):
+            self.crop = [self.crop]
+        else:
+            self.crop = list(self.crop)
+        if len(self.crop) == 1:
+            self.crop = [self.crop[0], (1 if self.crop[0] < 1 else size) - self.crop[0]] * dim
+        if len(self.crop) == dim:
+            crop = self.crop
+            self.crop = []
+            for value in crop:
+                self.crop += [value, (1 if value < 1 else size) - value]
+        assert len(self.crop) == 2 * dim, 'Crop values has to be ' + \
+            'a single value or a sequence of values of the same ' + \
+            'size as number of dimensions of the data or twice of that.'
+        pre_index = [Ellipsis]
+        post_index = [slice(None, None, None) for _ in range(dim)]
+        for i in range(0, 2 * dim, 2):
+            j = i // 2
+            size = data.size(2 + j)
+            crop_min, crop_max = self.crop[i:i + 2]
+            if crop_max < 1:
+                crop_min, crop_max = crop_min * size, crop_max * size
+            crop_min, crop_max = max(0, int(crop_min)), min(size, int(crop_max))
+            dim_index = post_index.copy()
+            dim_index[j] = slice(crop_min, crop_max, None)
+            data = data[pre_index + dim_index]
+        return data
+
+    def prep_latents(self, latents):
+        if self.full_sampling:
+            lerp = utils.slerp
+            if self.use_dlatent:
+                lerp = utils.lerp
+            latents_a, latents_b = latents[:self.batch_size], latents[self.batch_size:]
+            latents = lerp(
+                latents_a,
+                latents_b,
+                torch.rand(
+                    latents_a.size()[:-1],
+                    dtype=latents_a.dtype,
+                    device=latents_a.device
+                ).unsqueeze(-1)
+            )
+        return torch.cat([latents, latents + self.epsilon], dim=0)
+
+    def __call__(self, *args, **kwargs):
+        return self.evaluate(*args, **kwargs)
+
+    def evaluate(self, verbose=True):
+        """
+        Evaluate the PPL.
+        Arguments:
+            verbose (bool): Write progress to stdout.
+                Default value is True.
+        Returns:
+            ppl (float): Metric value.
+        """
+        distances = []
+        batch_size = self.batch_size
+        if self.full_sampling:
+            batch_size = 2 * batch_size
+
+        if verbose:
+            progress = utils.ProgressWriter(np.ceil(self.num_samples / self.batch_size))
+            progress.write('PPL: Evaluating metric...', step=False)
+
+        for _ in range(0, self.num_samples, self.batch_size):
+            utils.unwrap_module(self.G_synthesis).static_noise()
+
+            latents, latent_labels = self.prior_generator(batch_size=batch_size)
+            if latent_labels is not None and self.full_sampling:
+                # Labels should be the same for the first and second half of latents
+                latent_labels = latent_labels.view(2, -1)[0].repeat(2)
+
+            if self.use_dlatent:
+                with torch.no_grad():
+                    dlatents = self.G_mapping(latents=latents, labels=latent_labels)
+                dlatents = self.prep_latents(dlatents)
+            else:
+                latents = self.prep_latents(latents)
+                with torch.no_grad():
+                    dlatents = self.G_mapping(latents=latents, labels=latent_labels)
+
+            dlatents = dlatents.unsqueeze(1).repeat(1, len(utils.unwrap_module(self.G_synthesis)), 1)
+
+            with torch.no_grad():
+                output = self.G_synthesis(dlatents)
+
+            output = self.crop_data(output)
+            output = self._scale_for_lpips(output)
+
+            output_a, output_b = output[:self.batch_size], output[self.batch_size:]
+
+            with torch.no_grad():
+                dist = self.lpips_model(output_a, output_b)
+
+            distances.append(dist.cpu() * (1 / self.epsilon ** 2))
+
+            if verbose:
+                progress.step()
+
+        if verbose:
+            progress.write('PPL: Evaluated!', step=False)
+            progress.close()
+
+        distances = torch.cat(distances, dim=0).numpy()
+        lo = np.percentile(distances, 1, interpolation='lower')
+        hi = np.percentile(distances, 99, interpolation='higher')
+        filtered_distances = np.extract(np.logical_and(lo <= distances, distances <= hi), distances)
+        return float(np.mean(filtered_distances))

stylegan2/models.py

@@ -0,0 +1,1230 @@
+import copy
+from collections import OrderedDict
+import numpy as np
+import torch
+from torch import nn
+
+from . import modules, utils
+
+
+class _BaseModel(nn.Module):
+    """
+    Adds some base functionality to models that inherit this class.
+    """
+
+    def __init__(self):
+        super(_BaseModel, self).__setattr__('kwargs', {})
+        super(_BaseModel, self).__setattr__('_defaults', {})
+        super(_BaseModel, self).__init__()
+
+    def _update_kwargs(self, **kwargs):
+        """
+        Update the current keyword arguments. Overrides any
+        default values set.
+        Arguments:
+            **kwargs: Keyword arguments
+        """
+        self.kwargs.update(**kwargs)
+
+    def _update_default_kwargs(self, **defaults):
+        """
+        Update the default values for keyword arguments.
+        Arguments:
+            **defaults: Keyword arguments
+        """
+        self._defaults.update(**defaults)
+
+    def __getattr__(self, name):
+        """
+        Try to get the keyword argument for this attribute.
+        If no keyword argument of this name exists, try to
+        get the attribute directly from this object instead.
+        Arguments:
+            name (str): Name of keyword argument or attribute.
+        Returns:
+            value
+        """
+        try:
+            return self.__getattribute__('kwargs')[name]
+        except KeyError:
+            try:
+                return self.__getattribute__('_defaults')[name]
+            except KeyError:
+                return super(_BaseModel, self).__getattr__(name)
+
+    def __setattr__(self, name, value):
+        """
+        Try to set the keyword argument for this attribute.
+        If no keyword argument of this name exists, set
+        the attribute directly for this object instead.
+        Arguments:
+            name (str): Name of keyword argument or attribute.
+            value
+        """
+        if name != '__dict__' and (name in self.kwargs or name in self._defaults):
+            self.kwargs[name] = value
+        else:
+            super(_BaseModel, self).__setattr__(name, value)
+
+    def __delattr__(self, name):
+        """
+        Try to delete the keyword argument for this attribute.
+        If no keyword argument of this name exists, delete
+        the attribute of this object instead.
+        Arguments:
+            name (str): Name of keyword argument or attribute.
+        """
+        deleted = False
+        if name in self.kwargs:
+            del self.kwargs[name]
+            deleted = True
+        if name in self._defaults:
+            del self._defaults[name]
+            deleted = True
+        if not deleted:
+            super(_BaseModel, self).__delattr__(name)
+
+    def clone(self):
+        """
+        Create a copy of this model.
+        Returns:
+            model_copy (nn.Module)
+        """
+        return copy.deepcopy(self)
+
+    def _get_state_dict(self):
+        """
+        Delegate function for getting the state dict.
+        Should be overridden if state dict has to be
+        fetched in abnormal way.
+        """
+        return self.state_dict()
+
+    def _set_state_dict(self, state_dict):
+        """
+        Delegate function for loading the state dict.
+        Should be overridden if state dict has to be
+        loaded in abnormal way.
+        """
+        self.load_state_dict(state_dict)
+
+    def _serialize(self, half=False):
+        """
+        Turn model arguments and weights into
+        a dict that can safely be pickled and unpickled.
+        Arguments:
+            half (bool): Save weights in half precision.
+                Default value is False.
+        """
+        state_dict = self._get_state_dict()
+        for key in state_dict.keys():
+            values = state_dict[key].cpu()
+            if torch.is_floating_point(values):
+                if half:
+                    values = values.half()
+                else:
+                    values = values.float()
+            state_dict[key] = values
+        return {
+            'name': self.__class__.__name__,
+            'kwargs': self.kwargs,
+            'state_dict': state_dict
+        }
+
+    @classmethod
+    def load(cls, fpath, map_location='cpu'):
+        """
+        Load a model of this class.
+        Arguments:
+            fpath (str): File path of saved model.
+            map_location (str, int, torch.device): Weights and
+                buffers will be loaded into this device.
+                Default value is 'cpu'.
+        """
+        model = load(fpath, map_location=map_location)
+        assert isinstance(model, cls), 'Trying to load a `{}` '.format(type(model)) + \
+            'model from {}.load()'.format(cls.__name__)
+        return model
+
+    def save(self, fpath, half=False):
+        """
+        Save this model.
+        Arguments:
+            fpath (str): File path of save location.
+            half (bool): Save weights in half precision.
+                Default value is False.
+        """
+        torch.save(self._serialize(half=half), fpath)
+
+
+def _deserialize(state):
+    """
+    Load a model from its serialized state.
+    Arguments:
+        state (dict)
+    Returns:
+        model (nn.Module): Model that inherits `_BaseModel`.
+    """
+    state = state.copy()
+    name = state.pop('name')
+    if name not in globals():
+        raise NameError('Class {} is not defined.'.format(state['name']))
+    kwargs = state.pop('kwargs')
+    state_dict = state.pop('state_dict')
+    # Assume every other entry in the state is a serialized
+    # keyword argument.
+    for key in list(state.keys()):
+        kwargs[key] = _deserialize(state.pop(key))
+    model = globals()[name](**kwargs)
+    model._set_state_dict(state_dict)
+    return model
+
+
+def load(fpath, map_location='cpu'):
+    """
+    Load a model.
+    Arguments:
+        fpath (str): File path of saved model.
+        map_location (str, int, torch.device): Weights and
+            buffers will be loaded into this device.
+            Default value is 'cpu'.
+    Returns:
+        model (nn.Module): Model that inherits `_BaseModel`.
+    """
+    if map_location is not None:
+        map_location = torch.device(map_location)
+    return _deserialize(torch.load(fpath, map_location=map_location))
+
+
+def save(model, fpath, half=False):
+    """
+    Save a model.
+    Arguments:
+        model (nn.Module): Wrapped or unwrapped module
+            that inherits `_BaseModel`.
+        fpath (str): File path of save location.
+        half (bool): Save weights in half precision.
+            Default value is False.
+    """
+    utils.unwrap_module(model).save(fpath, half=half)
+
+
+class Generator(_BaseModel):
+    """
+    A wrapper class for the latent mapping model
+    and synthesis (generator) model.
+    Keyword Arguments:
+        G_mapping (GeneratorMapping)
+        G_synthesis (GeneratorSynthesis)
+        dlatent_avg_beta (float): The beta value
+            of the exponential moving average
+            of the dlatents. This statistic
+            is used for truncation of dlatents.
+            Default value is 0.995
+    """
+
+    def __init__(self, *, G_mapping, G_synthesis, **kwargs):
+        super(Generator, self).__init__()
+        self._update_default_kwargs(
+            dlatent_avg_beta=0.995
+        )
+        self._update_kwargs(**kwargs)
+
+        assert isinstance(G_mapping, GeneratorMapping), \
+            '`G_mapping` has to be an instance of `model.GeneratorMapping`'
+        assert isinstance(G_synthesis, GeneratorSynthesis), \
+            '`G_synthesis` has to be an instance of `model.GeneratorSynthesis`'
+        self.G_mapping = G_mapping
+        self.G_synthesis = G_synthesis
+        self.register_buffer('dlatent_avg', torch.zeros(self.G_mapping.latent_size))
+        self.set_truncation()
+
+    @property
+    def latent_size(self):
+        return self.G_mapping.latent_size
+
+    @property
+    def label_size(self):
+        return self.G_mapping.label_size
+
+    def _get_state_dict(self):
+        state_dict = OrderedDict()
+        self._save_to_state_dict(destination=state_dict, prefix='', keep_vars=False)
+        return state_dict
+
+    def _set_state_dict(self, state_dict):
+        self.load_state_dict(state_dict, strict=False)
+
+    def _serialize(self, half=False):
+        state = super(Generator, self)._serialize(half=half)
+        for name in ['G_mapping', 'G_synthesis']:
+            state[name] = getattr(self, name)._serialize(half=half)
+        return state
+
+    def set_truncation(self, truncation_psi=None, truncation_cutoff=None):
+        """
+        Set the truncation of dlatents before they are passed to the
+        synthesis model.
+        Arguments:
+            truncation_psi (float): Beta value of linear interpolation between
+                the average dlatent and the current dlatent. 0 -> 100% average,
+                1 -> 0% average.
+            truncation_cutoff (int, optional): Truncation is only used up until
+                this affine layer index.
+        """
+        layer_psi = None
+        if truncation_psi is not None and truncation_psi != 1 and truncation_cutoff != 0:
+            layer_psi = torch.ones(len(self.G_synthesis))
+            if truncation_cutoff is None:
+                layer_psi *= truncation_psi
+            else:
+                layer_psi_mask = torch.arange(len(layer_psi)) < truncation_cutoff
+                layer_psi[layer_psi_mask] *= truncation_psi
+            layer_psi = layer_psi.view(1, -1, 1)
+            layer_psi = layer_psi.to(self.dlatent_avg)
+        self.register_buffer('layer_psi', layer_psi)
+
+    def random_noise(self):
+        """
+        Set noise of synthesis model to be random for every
+        input.
+        """
+        self.G_synthesis.random_noise()
+
+    def static_noise(self, trainable=False, noise_tensors=None):
+        """
+        Set up injected noise to be fixed (alternatively trainable).
+        Get the fixed noise tensors (or parameters).
+        Arguments:
+            trainable (bool): Make noise trainable and return
+                parameters instead of normal tensors.
+            noise_tensors (list, optional): List of tensors to use as static noise.
+                Has to be same length as number of noise injection layers.
+        Returns:
+            noise_tensors (list): List of the noise tensors (or parameters).
+        """
+        return self.G_synthesis.static_noise(trainable=trainable, noise_tensors=noise_tensors)
+
+    def __len__(self):
+        """
+        Get the number of affine (style) layers of the synthesis model.
+        """
+        return len(self.G_synthesis)
+
+    def truncate(self, dlatents):
+        """
+        Truncate the dlatents.
+        Arguments:
+            dlatents (torch.Tensor)
+        Returns:
+            truncated_dlatents (torch.Tensor)
+        """
+        if self.layer_psi is not None:
+            dlatents = utils.lerp(self.dlatent_avg, dlatents, self.layer_psi)
+        return dlatents
+
+    def forward(self,
+                latents=None,
+                labels=None,
+                dlatents=None,
+                return_dlatents=False,
+                mapping_grad=True,
+                latent_to_layer_idx=None):
+        """
+        Synthesize some data from latent inputs. The latents
+        can have an extra optional dimension, where latents
+        from this dimension will be distributed to the different
+        affine layers of the synthesis model. The distribution
+        is a index to index mapping if the amount of latents
+        is the same as the number of affine layers. Otherwise,
+        latents are distributed consecutively for a random
+        number of layers before the next latent is used for
+        some random amount of following layers. If the size
+        of this extra dimension is 1 or it does not exist,
+        the same latent is passed to every affine layer.
+
+        Latents are first mapped to disentangled latents (`dlatents`)
+        and are then optionally truncated (if model is in eval mode
+        and truncation options have been set.) Set up truncation by
+        calling `set_truncation()`.
+        Arguments:
+            latents (torch.Tensor): The latent values of shape
+                (batch_size, N, num_features) where N is an
+                optional dimension. This argument is not required
+                if `dlatents` is passed.
+            labels (optional): A sequence of labels, one for
+                each index in the batch dimension of the input.
+            dlatents (torch.Tensor, optional): Skip the latent
+                mapping model and feed these dlatents straight
+                to the synthesis model. The same type of distribution
+                to affine layers as is described in this function
+                description is also used for dlatents.
+                NOTE: Explicitly passing dlatents to this function
+                    will stop them from being truncated. If required,
+                    do this manually by calling the `truncate()` function
+                    of this model.
+            return_dlatents (bool): Return not only the synthesized
+                data, but also the dlatents. The dlatents tensor
+                will also have its `requires_grad` set to True
+                before being passed to the synthesis model for
+                use with pathlength regularization during training.
+                This requires training to be enabled (`thismodel.train()`).
+                Default value is False.
+            mapping_grad (bool): Let gradients be calculated when passing
+                latents through the latent mapping model. Should be
+                set to False when only optimising the synthesiser parameters.
+                Default value is True.
+            latent_to_layer_idx (list, tuple, optional): A manual mapping
+                of the latent vectors to the affine layers of this network.
+                Each position in this sequence maps the affine layer of the
+                same index to an index of the latents. The latents should
+                have a shape of (batch_size, N, num_features) and this argument
+                should be a list of the same length as number of affine layers
+                in this model (can be found by calling len(thismodel)) with values
+                in the range [0, N - 1]. Without this argument, latents are distributed
+                according to this function description.
+        """
+        # Keep track of number of latents for each batch index.
+        num_latents = 1
+
+        # Keep track of if dlatent truncation is enabled or disabled.
+        truncate = False
+
+        if dlatents is None:
+            # Calculate dlatents
+
+            # dlatent truncation enabled as dlatents were not explicitly given
+            truncate = True
+
+            assert latents is not None, 'Either the `latents` ' + \
+                'or the `dlatents` argument is required.'
+            if labels is not None:
+                if not torch.is_tensor(labels):
+                    labels = torch.tensor(labels, dtype=torch.int64)
+
+            # If latents are passed with the layer dimension we need
+            # to flatten it to shape (N, latent_size) before passing
+            # it to the latent mapping model.
+            if latents.dim() == 3:
+                num_latents = latents.size(1)
+                latents = latents.view(-1, latents.size(-1))
+                # Labels need to repeated for the extra dimension of latents.
+                if labels is not None:
+                    labels = labels.unsqueeze(1).repeat(1, num_latents).view(-1)
+
+            # Dont allow this operation to create a computation graph for
+            # backprop unless specified. This is useful for pathreg as it
+            # only regularizes the parameters of the synthesiser and not
+            # to latent mapping model.
+            with torch.set_grad_enabled(mapping_grad):
+                dlatents = self.G_mapping(latents=latents, labels=labels)
+        else:
+            if dlatents.dim() == 3:
+                num_latents = dlatents.size(1)
+
+        # Now we expand/repeat the number of latents per batch index until it is
+        # the same number as affine layers in our synthesis model.
+        dlatents = dlatents.view(-1, num_latents, dlatents.size(-1))
+        if num_latents == 1:
+            dlatents = dlatents.expand(
+                dlatents.size(0), len(self), dlatents.size(2))
+        elif num_latents != len(self):
+            assert dlatents.size(1) <= len(self), \
+                'More latents ({}) than number '.format(dlatents.size(1)) + \
+                'of generator layers ({}) received.'.format(len(self))
+            if not latent_to_layer_idx:
+                # Lets randomly distribute the latents to
+                # ranges of layers (each latent is assigned
+                # to a random number of consecutive layers).
+                cutoffs = np.random.choice(
+                    np.arange(1, len(self)),
+                    dlatents.size(1) - 1,
+                    replace=False
+                )
+                cutoffs = [0] + sorted(cutoffs.tolist()) + [len(self)]
+                dlatents = [
+                    dlatents[:, i].unsqueeze(1).expand(
+                        -1, cutoffs[i + 1] - cutoffs[i], dlatents.size(2))
+                    for i in range(dlatents.size(1))
+                ]
+                dlatents = torch.cat(dlatents, dim=1)
+            else:
+                # Assign latents as specified by argument
+                assert len(latent_to_layer_idx) == len(self), \
+                    'The latent index to layer index mapping does ' + \
+                    'not have the same number of elements ' + \
+                    '({}) as the number of '.format(len(latent_to_layer_idx)) + \
+                    'generator layers ({})'.format(len(self))
+                dlatents = dlatents[:, latent_to_layer_idx]
+
+        # Update moving average of dlatents when training
+        if self.training and self.dlatent_avg_beta != 1:
+            with torch.no_grad():
+                batch_dlatent_avg = dlatents[:, 0].mean(dim=0)
+                self.dlatent_avg = utils.lerp(
+                    batch_dlatent_avg, self.dlatent_avg, self.dlatent_avg_beta)
+
+        # Truncation is only applied when dlatents are not explicitly
+        # given and the model is in evaluation mode.
+        if truncate and not self.training:
+            dlatents = self.truncate(dlatents)
+
+        # One of the reasons we might want to return the dlatents is for
+        # pathreg, in which case the dlatents need to require gradients
+        # before being passed to the synthesiser. This should only be
+        # the case when the model is in training mode.
+        if return_dlatents and self.training:
+            dlatents.requires_grad_(True)
+
+        synth = self.G_synthesis(latents=dlatents)
+        if return_dlatents:
+            return synth, dlatents
+        return synth
+
+
+# Base class for the parameterized models. This is used as parent
+# class to reduce duplicate code and documentation for shared arguments.
+class _BaseParameterizedModel(_BaseModel):
+    """
+        activation (str, callable, nn.Module): The non-linear
+            activation function to use.
+            Default value is leaky relu with a slope of 0.2.
+        lr_mul (float): The learning rate multiplier for this
+            model. When loading weights of previously trained
+            networks, this value has to be the same as when
+            the network was trained for the outputs to not
+            change (as this is used to scale the weights).
+            Default value depends on model type and can
+            be found in the original paper for StyleGAN.
+        weight_scale (bool): Use weight scaling for
+            equalized learning rate. Default value
+            is True.
+        eps (float): Epsilon value added for numerical stability.
+            Default value is 1e-8."""
+
+    def __init__(self, **kwargs):
+        super(_BaseParameterizedModel, self).__init__()
+        self._update_default_kwargs(
+             activation='lrelu:0.2',
+             lr_mul=1,
+             weight_scale=True,
+             eps=1e-8
+        )
+        self._update_kwargs(**kwargs)
+
+
+class GeneratorMapping(_BaseParameterizedModel):
+    """
+    Latent mapping model, handles the
+    transformation of latents into disentangled
+    latents.
+    Keyword Arguments:
+        latent_size (int): The size of the latent vectors.
+            This will also be the size of the disentangled
+            latent vectors.
+            Default value is 512.
+        label_size (int, optional): The number of different
+            possible labels. Use for label conditioning of
+            the GAN. Unused by default.
+        out_size (int, optional): The size of the disentangled
+            latents output by this model. If not specified,
+            the outputs will have the same size as the input
+            latents.
+        num_layers (int): Number of dense layers in this
+            model. Default value is 8.
+        hidden (int, optional): Number of hidden features of layers.
+            If unspecified, this is the same size as the latents.
+        normalize_input (bool): Normalize the input of this
+            model. Default value is True."""
+    __doc__ += _BaseParameterizedModel.__doc__
+
+    def __init__(self, **kwargs):
+        super(GeneratorMapping, self).__init__()
+        self._update_default_kwargs(
+            latent_size=512,
+            label_size=0,
+            out_size=None,
+            num_layers=8,
+            hidden=None,
+            normalize_input=True,
+            lr_mul=0.01,
+        )
+        self._update_kwargs(**kwargs)
+
+        # Find in and out features of first dense layer
+        in_features = self.latent_size
+        out_features = self.hidden or self.latent_size
+
+        # Each class label has its own embedded vector representation.
+        self.embedding = None
+        if self.label_size:
+            self.embedding = nn.Embedding(self.label_size, self.latent_size)
+            # The input is now the latents concatenated with
+            # the label embeddings.
+            in_features += self.latent_size
+        dense_layers = []
+        for i in range(self.num_layers):
+            if i == self.num_layers - 1:
+                # Set out features for last dense layer
+                out_features = self.out_size or self.latent_size
+            dense_layers.append(
+                modules.BiasActivationWrapper(
+                    layer=modules.DenseLayer(
+                        in_features=in_features,
+                        out_features=out_features,
+                        lr_mul=self.lr_mul,
+                        weight_scale=self.weight_scale,
+                        gain=1
+                    ),
+                    features=out_features,
+                    use_bias=True,
+                    activation=self.activation,
+                    bias_init=0,
+                    lr_mul=self.lr_mul,
+                    weight_scale=self.weight_scale
+                )
+            )
+            in_features = out_features
+        self.main = nn.Sequential(*dense_layers)
+
+    def forward(self, latents, labels=None):
+        """
+        Get the disentangled latents from the input latents
+        and optional labels.
+        Arguments:
+            latents (torch.Tensor): Tensor of shape (batch_size, latent_size).
+            labels (torch.Tensor, optional): Labels for conditioning of latents
+                if there are any.
+        Returns:
+            dlatents (torch.Tensor): Disentangled latents of same shape as
+                `latents` argument.
+        """
+        assert latents.dim() == 2 and latents.size(-1) == self.latent_size, \
+            'Incorrect input shape. Should be ' + \
+            '(batch_size, {}) '.format(self.latent_size) + \
+            'but received {}'.format(tuple(latents.size()))
+        x = latents
+        if labels is not None:
+            assert self.embedding is not None, \
+                'No embedding layer found, please ' + \
+                'specify the number of possible labels ' + \
+                'in the constructor of this class if ' + \
+                'using labels.'
+            assert len(labels) == len(latents), \
+                'Received different number of labels ' + \
+                '({}) and latents ({}).'.format(len(labels), len(latents))
+            if not torch.is_tensor(labels):
+                labels = torch.tensor(labels, dtype=torch.int64)
+            assert labels.dtype == torch.int64, \
+                'Labels should be integer values ' + \
+                'of dtype torch.in64 (long)'
+            y = self.embedding(labels)
+            x = torch.cat([x, y], dim=-1)
+        else:
+            assert self.embedding is None, 'Missing input labels.'
+        if self.normalize_input:
+            x = x * torch.rsqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
+        return self.main(x)
+
+
+# Base class for the synthesising and discriminating models. This is used as parent
+# class to reduce duplicate code and documentation for shared arguments.
+class _BaseAdverserialModel(_BaseParameterizedModel):
+    """
+        data_channels (int): Number of channels of the data.
+            Default value is 3.
+        base_shape (list, tuple): This is the shape of the feature
+            activations when it is most compact and still has the
+            same number of dims as the data. This is one of the
+            arguments that controls what shape the data will be.
+            the value of each size in the shape is going to double
+            in size for number of `channels` - 1.
+            Example:
+                `data_channels=3`
+                `base_shape=(4, 2)`
+                and 9 `channels` in total will give us a shape of
+                (3, 4 * 2^(9 - 1), 2 * 2^(9 - 1)) which is the
+                same as (3, 1024, 512).
+            Default value is (4, 4).
+        channels (int, list, optional): The channels of each block
+            of layers. If int, this many channel values will be
+            created with sensible default values optimal for image
+            synthesis. If list, the number of blocks in this model
+            will be the same as the number of channels in the list.
+            Default value is the int value 9 which will create the
+            following channels: [32, 32, 64, 128, 256, 512, 512, 512, 512].
+            These are the channel values used in the stylegan2 paper for
+            their FFHQ-trained face generation network.
+            If channels is given as a list it should be in the order:
+                Generator: last layer -> first layer
+                Discriminator: first layer -> last layer
+        resnet (bool): Use resnet connections.
+            Defaults:
+                Generator: False
+                Discriminator: True
+        skip (bool): Use skip connections for data.
+            Defaults:
+                Generator: True
+                Discriminator: False
+        fused_resample (bool): Fuse any up- or downsampling that
+            is paired with a convolutional layer into a strided
+            convolution (transposed if upsampling was used).
+            Default value is True.
+        conv_resample_mode (str): The resample mode of up- or
+            downsampling layers. If `fused_resample=True` only
+            'FIR' and 'none' can be used. Else, 'FIR' or anything
+            that can be passed to torch.nn.functional.interpolate
+            is a valid mode (and 'max' but only for downsampling
+            operations). Default value is 'FIR'.
+        conv_filter (int, list): The filter to use if
+            `conv_resample_mode='FIR'`. If int, a low
+            pass filter of this size will be used. If list,
+            the filter is explicitly specified. If the filter
+            is of a single dimension it will be expanded to
+            the number of dimensions of the data. Default
+            value is a low pass filter of [1, 3, 3, 1].
+        skip_resample_mode (str): If `skip=True`, this
+            mode is used for the resamplings of skip
+            connections of different sizes. Same possible
+            values as `conv_filter` (except 'none', which
+            can not be used). Default value is 'FIR'.
+        skip_filter (int, list): Same description as
+            `conv_filter` but for skip connections.
+            Only used if `skip_resample_mode='FIR'` and
+            `skip=True`. Default value is a low pass
+            filter of [1, 3, 3, 1].
+        kernel_size (int): The size of the convolutional kernels.
+            Default value is 3.
+        conv_pad_mode (str): The padding mode for convolutional
+            layers. Has to be one of 'constant', 'reflect',
+            'replicate' or 'circular'. Default value is
+            'constant'.
+        conv_pad_constant (float): The value to use for conv
+            padding if `conv_pad_mode='constant'`. Default
+            value is 0.
+        filter_pad_mode (str): The padding mode for FIR
+            filters. Same possible values as `conv_pad_mode`.
+            Default value is 'constant'.
+        filter_pad_constant (float): The value to use for FIR
+            padding if `filter_pad_mode='constant'`. Default
+            value is 0.
+        pad_once (bool): If FIR filter is used in conjunction with a
+            conv layer, do all the padding for both convolution and
+            FIR in the FIR layer instead of once per layer.
+            Default value is True.
+        conv_block_size (int): The number of conv layers in
+            each conv block. Default value is 2."""
+    __doc__ += _BaseParameterizedModel.__doc__
+
+    def __init__(self, **kwargs):
+        super(_BaseAdverserialModel, self).__init__()
+        self._update_default_kwargs(
+            data_channels=3,
+            base_shape=(4, 4),
+            channels=9,
+            resnet=False,
+            skip=False,
+            fused_resample=True,
+            conv_resample_mode='FIR',
+            conv_filter=[1, 3, 3, 1],
+            skip_resample_mode='FIR',
+            skip_filter=[1, 3, 3, 1],
+            kernel_size=3,
+            conv_pad_mode='constant',
+            conv_pad_constant=0,
+            filter_pad_mode='constant',
+            filter_pad_constant=0,
+            pad_once=True,
+            conv_block_size=2,
+        )
+        self._update_kwargs(**kwargs)
+
+        self.dim = len(self.base_shape)
+        assert 1 <= self.dim <= 3, '`base_shape` can only have 1, 2 or 3 dimensions.'
+        if isinstance(self.channels, int):
+            # Create the specified number of channel values with sensible
+            # sizes (these values do well for image synthesis).
+            num_channels = self.channels
+            self.channels = [min(32 * 2 ** i, 512) for i in range(min(8, num_channels))]
+            if len(self.channels) < num_channels:
+                self.channels = [32] * (num_channels - len(self.channels)) + self.channels
+
+
+class GeneratorSynthesis(_BaseAdverserialModel):
+    """
+    The synthesis model that takes latents and synthesises
+    some data.
+    Keyword Arguments:
+        latent_size (int): The size of the latent vectors.
+            This will also be the size of the disentangled
+            latent vectors.
+            Default value is 512.
+        demodulate (bool): Normalize feature outputs from conv
+            layers. Default value is True.
+        modulate_data_out (bool): Apply style to the data output
+            layers. These layers are projections of the feature
+            maps into the space of the data. Default value is True.
+        noise (bool): Add noise after each conv style layer.
+            Default value is True."""
+    __doc__ += _BaseAdverserialModel.__doc__
+
+    def __init__(self, **kwargs):
+        super(GeneratorSynthesis, self).__init__()
+        self._update_default_kwargs(
+            latent_size=512,
+            demodulate=True,
+            modulate_data_out=True,
+            noise=True,
+            resnet=False,
+            skip=True
+        )
+        self._update_kwargs(**kwargs)
+
+        # The constant input of the model has no activations
+        # normalization, it is just passed straight to the first
+        # layer of the model.
+        self.const = torch.nn.Parameter(
+            torch.empty(self.channels[-1], *self.base_shape).normal_()
+        )
+        conv_block_kwargs = dict(
+            latent_size=self.latent_size,
+            demodulate=self.demodulate,
+            resnet=self.resnet,
+            up=True,
+            num_layers=self.conv_block_size,
+            filter=self.conv_filter,
+            activation=self.activation,
+            mode=self.conv_resample_mode,
+            fused=self.fused_resample,
+            kernel_size=self.kernel_size,
+            pad_mode=self.conv_pad_mode,
+            pad_constant=self.conv_pad_constant,
+            filter_pad_mode=self.filter_pad_mode,
+            filter_pad_constant=self.filter_pad_constant,
+            pad_once=self.pad_once,
+            noise=self.noise,
+            lr_mul=self.lr_mul,
+            weight_scale=self.weight_scale,
+            gain=1,
+            dim=self.dim,
+            eps=self.eps
+        )
+        self.conv_blocks = nn.ModuleList()
+
+        # The first convolutional layer is slightly different
+        # from the following convolutional blocks but can still
+        # be represented as a convolutional block if we change
+        # some of its arguments.
+        self.conv_blocks.append(
+            modules.GeneratorConvBlock(
+                **{
+                    **conv_block_kwargs,
+                    'in_channels': self.channels[-1],
+                    'out_channels': self.channels[-1],
+                    'resnet': False,
+                    'up': False,
+                    'num_layers': 1
+                }
+            )
+        )
+
+        # The rest of the convolutional blocks all look the same
+        # except for number of input and output channels
+        for i in range(1, len(self.channels)):
+            self.conv_blocks.append(
+                modules.GeneratorConvBlock(
+                    in_channels=self.channels[-i],
+                    out_channels=self.channels[-i - 1],
+                    **conv_block_kwargs
+                )
+            )
+
+        # If not using the skip architecture, only one
+        # layer will project the feature maps into
+        # the space of the data (from the activations of
+        # the last convolutional block). If using the skip
+        # architecture, every block will have its
+        # own projection layer instead.
+        self.to_data_layers = nn.ModuleList()
+        for i in range(1, len(self.channels) + 1):
+            to_data = None
+            if i == len(self.channels) or self.skip:
+                to_data = modules.BiasActivationWrapper(
+                    layer=modules.ConvLayer(
+                        **{
+                            **conv_block_kwargs,
+                            'in_channels': self.channels[-i],
+                            'out_channels': self.data_channels,
+                            'modulate': self.modulate_data_out,
+                            'demodulate': False,
+                            'kernel_size': 1
+                        }
+                    ),
+                    **{
+                        **conv_block_kwargs,
+                        'features': self.data_channels,
+                        'use_bias': True,
+                        'activation': 'linear',
+                        'bias_init': 0
+                    }
+                )
+            self.to_data_layers.append(to_data)
+
+        # When the skip architecture is used we need to
+        # upsample data outputs of previous convolutional
+        # blocks so that it can be added to the data output
+        # of the current convolutional block.
+        self.upsample = None
+        if self.skip:
+            self.upsample = modules.Upsample(
+                mode=self.skip_resample_mode,
+                filter=self.skip_filter,
+                filter_pad_mode=self.filter_pad_mode,
+                filter_pad_constant=self.filter_pad_constant,
+                gain=1,
+                dim=self.dim
+            )
+
+        # Calculate the number of latents required
+        # in the input.
+        self._num_latents = 1 + self.conv_block_size * (len(self.channels) - 1)
+        # Only the final data output layer uses
+        # its own latent input when being modulated.
+        # The other data output layers recycles latents
+        # from the next convolutional block.
+        if self.modulate_data_out:
+            self._num_latents += 1
+
+    def __len__(self):
+        """
+        Get the number of affine (style) layers of this model.
+        """
+        return self._num_latents
+
+    def random_noise(self):
+        """
+        Set injected noise to be random for each new input.
+        """
+        for module in self.modules():
+            if isinstance(module, modules.NoiseInjectionWrapper):
+                module.random_noise()
+
+    def static_noise(self, trainable=False, noise_tensors=None):
+        """
+        Set up injected noise to be fixed (alternatively trainable).
+        Get the fixed noise tensors (or parameters).
+        Arguments:
+            trainable (bool): Make noise trainable and return
+                parameters instead of normal tensors.
+            noise_tensors (list, optional): List of tensors to use as static noise.
+                Has to be same length as number of noise injection layers.
+        Returns:
+            noise_tensors (list): List of the noise tensors (or parameters).
+        """
+        rtn_tensors = []
+
+        if not self.noise:
+            return rtn_tensors
+
+        for module in self.modules():
+            if isinstance(module, modules.NoiseInjectionWrapper):
+                has_noise_shape = module.has_noise_shape()
+                device = module.weight.device
+                dtype = module.weight.dtype
+                break
+
+        # If noise layers dont have the shape that the noise should be
+        # we first need to pass some data through the network once for
+        # these layers to record the shape. To create noise tensors
+        # we need to know what size they should be.
+        if not has_noise_shape:
+            with torch.no_grad():
+                self(torch.zeros(
+                    1, len(self), self.latent_size, device=device, dtype=dtype))
+
+        i = 0
+        for block in self.conv_blocks:
+            for layer in block.conv_block:
+                for module in layer.modules():
+                    if isinstance(module, modules.NoiseInjectionWrapper):
+                        noise_tensor = None
+                        if noise_tensors is not None:
+                            if i < len(noise_tensors):
+                                noise_tensor = noise_tensors[i]
+                                i += 1
+                            else:
+                                rtn_tensors.append(None)
+                                continue
+                        rtn_tensors.append(
+                            module.static_noise(trainable=trainable, noise_tensor=noise_tensor))
+
+        if noise_tensors is not None:
+            assert len(rtn_tensors) == len(noise_tensors), \
+                'Got a list of {} '.format(len(noise_tensors)) + \
+                'noise tensors but there are ' + \
+                '{} noise layers in this model'.format(len(rtn_tensors))
+
+        return rtn_tensors
+
+    def forward(self, latents):
+        """
+        Synthesise some data from input latents.
+        Arguments:
+            latents (torch.Tensor): Latent vectors of shape
+                (batch_size, num_affine_layers, latent_size)
+                where num_affine_layers is the value returned
+                by __len__() of this class.
+        Returns:
+            synthesised (torch.Tensor): Synthesised data.
+        """
+        assert latents.dim() == 3 and latents.size(1) == len(self), \
+            'Input mismatch, expected latents of shape ' + \
+            '(batch_size, {}, latent_size) '.format(len(self)) + \
+            'but got {}.'.format(tuple(latents.size()))
+        # Declare our feature activations variable
+        # and give it the value of our const parameter with
+        # an added batch dimension.
+        x = self.const.unsqueeze(0)
+        # Declare our data (output) variable
+        y = None
+        # Start counting style layers used. This is used for specifying
+        # which latents should be passed to the current block in the loop.
+        layer_idx = 0
+        for block, to_data in zip(self.conv_blocks, self.to_data_layers):
+            # Get the latents for the style layers in this block.
+            block_latents = latents[:, layer_idx:layer_idx + len(block)]
+
+            x = block(input=x, latents=block_latents)
+
+            layer_idx += len(block)
+
+            # Upsample the data output of the previous block to fit
+            # the data output size of this block so that they can
+            # be added together. Only performed for 'skip' architectures.
+            if self.upsample is not None and layer_idx < len(self):
+                if y is not None:
+                    y = self.upsample(y)
+
+            # Combine the data output of this block with any previous
+            # blocks outputs if using 'skip' architecture, else only
+            # perform this operation for the very last block outputs.
+            if to_data is not None:
+                t = to_data(input=x, latent=latents[:, layer_idx])
+                y = t if y is None else y + t
+        return y
+
+
+class Discriminator(_BaseAdverserialModel):
+    """
+    The discriminator scores data inputs.
+    Keyword Arguments:
+        label_size (int, optional): The number of different
+            possible labels. Use for label conditioning of
+            the GAN. The discriminator will calculate scores
+            for each possible label and only returns the score
+            from the label passed with the input data. If no
+            labels are used, only one score is calculated.
+            Disabled by default.
+        mbstd_group_size (int): Group size for minibatch std
+            before the final conv layer. A value of 0 indicates
+            not to use minibatch std, and a value of -1 indicates
+            that the group should be over the entire batch.
+            This is used for increasing variety of the outputs of
+            the generator. Default value is 4.
+            NOTE: Scores for the same data may vary depending
+                on batch size when using a value of -1.
+            NOTE: If a value > 0 is given, every input batch
+                must have a size evenly divisible by this value.
+        dense_hidden (int, optional): The number of hidden features
+            of the first dense layer. By default, this is the same as
+            the number of channels in the final conv layer."""
+    __doc__ += _BaseAdverserialModel.__doc__
+
+    def __init__(self, **kwargs):
+        super(Discriminator, self).__init__()
+        self._update_default_kwargs(
+            label_size=0,
+            mbstd_group_size=4,
+            dense_hidden=None,
+            resnet=True,
+            skip=False
+        )
+        self._update_kwargs(**kwargs)
+
+        conv_block_kwargs = dict(
+            resnet=self.resnet,
+            down=True,
+            num_layers=self.conv_block_size,
+            filter=self.conv_filter,
+            activation=self.activation,
+            mode=self.conv_resample_mode,
+            fused=self.fused_resample,
+            kernel_size=self.kernel_size,
+            pad_mode=self.conv_pad_mode,
+            pad_constant=self.conv_pad_constant,
+            filter_pad_mode=self.filter_pad_mode,
+            filter_pad_constant=self.filter_pad_constant,
+            pad_once=self.pad_once,
+            noise=False,
+            lr_mul=self.lr_mul,
+            weight_scale=self.weight_scale,
+            gain=1,
+            dim=self.dim,
+            eps=self.eps
+        )
+        self.conv_blocks = nn.ModuleList()
+
+        # All but the last of the convolutional blocks look the same
+        # except for number of input and output channels
+        for i in range(len(self.channels) - 1):
+            self.conv_blocks.append(
+                modules.DiscriminatorConvBlock(
+                    in_channels=self.channels[i],
+                    out_channels=self.channels[i + 1],
+                    **conv_block_kwargs
+                )
+            )
+
+        # The final convolutional layer is slightly different
+        # from the previous convolutional blocks but can still
+        # be represented as a convolutional block if we change
+        # some of its arguments and optionally add a minibatch
+        # std layer before it.
+        final_conv_block = []
+        if self.mbstd_group_size:
+            final_conv_block.append(
+                modules.MinibatchStd(
+                    group_size=self.mbstd_group_size,
+                    eps=self.eps
+                )
+            )
+        final_conv_block.append(
+            modules.DiscriminatorConvBlock(
+                **{
+                    **conv_block_kwargs,
+                    'in_channels': self.channels[-1] + (1 if self.mbstd_group_size else 0),
+                    'out_channels': self.channels[-1],
+                    'resnet': False,
+                    'down': False,
+                    'num_layers': 1
+                },
+            )
+        )
+        self.conv_blocks.append(nn.Sequential(*final_conv_block))
+
+        # If not using the skip architecture, only one
+        # layer will project the data into feature maps.
+        # This would be performed only for the input data at
+        # the first block.
+        # If using the skip architecture, every block will
+        # have its own projection layer instead.
+        self.from_data_layers = nn.ModuleList()
+        for i in range(len(self.channels)):
+            from_data = None
+            if i == 0 or self.skip:
+                from_data = modules.BiasActivationWrapper(
+                    layer=modules.ConvLayer(
+                        **{
+                            **conv_block_kwargs,
+                            'in_channels': self.data_channels,
+                            'out_channels': self.channels[i],
+                            'modulate': False,
+                            'demodulate': False,
+                            'kernel_size': 1
+                        }
+                    ),
+                    **{
+                        **conv_block_kwargs,
+                        'features': self.channels[i],
+                        'use_bias': True,
+                        'activation': self.activation,
+                        'bias_init': 0
+                    }
+                )
+            self.from_data_layers.append(from_data)
+
+        # When the skip architecture is used we need to
+        # downsample the data input so that it has the same
+        # size as the feature maps of each block so that it
+        # can be projected and added to these feature maps.
+        self.downsample = None
+        if self.skip:
+            self.downsample = modules.Downsample(
+                mode=self.skip_resample_mode,
+                filter=self.skip_filter,
+                filter_pad_mode=self.filter_pad_mode,
+                filter_pad_constant=self.filter_pad_constant,
+                gain=1,
+                dim=self.dim
+            )
+
+        # The final layers are two dense layers that maps
+        # the features into score logits. If labels are
+        # used, we instead output one score for each possible
+        # class of the labels and then return the score for the
+        # labeled class.
+        dense_layers = []
+        in_features = self.channels[-1] * np.prod(self.base_shape)
+        out_features = self.dense_hidden or self.channels[-1]
+        activation = self.activation
+        for _ in range(2):
+            dense_layers.append(
+                modules.BiasActivationWrapper(
+                    layer=modules.DenseLayer(
+                        in_features=in_features,
+                        out_features=out_features,
+                        lr_mul=self.lr_mul,
+                        weight_scale=self.weight_scale,
+                        gain=1,
+                    ),
+                    features=out_features,
+                    activation=activation,
+                    use_bias=True,
+                    bias_init=0,
+                    lr_mul=self.lr_mul,
+                    weight_scale=self.weight_scale
+                )
+            )
+            in_features = out_features
+            out_features = max(1, self.label_size)
+            activation = 'linear'
+        self.dense = nn.Sequential(*dense_layers)
+
+    def forward(self, input, labels=None):
+        """
+        Takes some data and optionally its labels and
+        produces one score logit per data input.
+        Arguments:
+            input (torch.Tensor)
+            labels (torch.Tensor, list, optional)
+        Returns:
+            score_logits (torch.Tensor)
+        """
+        # Declare our feature activations variable.
+        x = None
+        # Declare our data (input) variable
+        y = input
+        for i, (block, from_data) in enumerate(zip(self.conv_blocks, self.from_data_layers)):
+            # Combine the data input of this block with any previous
+            # block output if using 'skip' architecture, else only
+            # perform this operation as a way to create inputs for
+            # the first block.
+            if from_data is not None:
+                t = from_data(y)
+                x = t if x is None else x + t
+
+            x = block(input=x)
+
+            # Downsample the data input of this block to fit
+            # the feature size of the output of this block so that they can
+            # be added together. Only performed for 'skip' architectures.
+            if self.downsample is not None and i != len(self.conv_blocks) - 1:
+                y = self.downsample(y)
+        # Calculate scores
+        x = x.view(x.size(0), -1)
+        x = self.dense(x)
+        if labels is not None:
+            # Use advanced indexing to fetch only the score of the
+            # class labels.
+            x = x[torch.arange(x.size(0)), labels].unsqueeze(-1)
+        return x

stylegan2/modules.py

@@ -0,0 +1,1601 @@
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+def get_activation(activation):
+    """
+    Get the module for a specific activation function and its gain if
+    it can be calculated.
+    Arguments:
+        activation (str, callable, nn.Module): String representing the activation.
+    Returns:
+        activation_module (torch.nn.Module): The module representing
+            the activation function.
+        gain (float): The gain value. Defaults to 1 if it can not be calculated.
+    """
+    if isinstance(activation, nn.Module) or callable(activation):
+        return activation, 1.
+    if isinstance(activation, str):
+        activation = activation.lower()
+    if activation in [None, 'linear']:
+        return nn.Identity(), 1.
+    lrelu_strings = ('leaky', 'leakyrely', 'leaky_relu', 'leaky relu', 'lrelu')
+    if activation.startswith(lrelu_strings):
+        for l_s in lrelu_strings:
+            activation = activation.replace(l_s, '')
+        slope = ''.join(
+            char for char in activation if char.isdigit() or char == '.')
+        slope = float(slope) if slope else 0.01
+        return nn.LeakyReLU(slope), np.sqrt(2)  # close enough to true gain
+    elif activation.startswith('swish'):
+        return Swish(affine=activation != 'swish'), np.sqrt(2)
+    elif activation in ['relu']:
+        return nn.ReLU(), np.sqrt(2)
+    elif activation in ['elu']:
+        return nn.ELU(), 1.
+    elif activation in ['prelu']:
+        return nn.PReLU(), np.sqrt(2)
+    elif activation in ['rrelu', 'randomrelu']:
+        return nn.RReLU(), np.sqrt(2)
+    elif activation in ['selu']:
+        return nn.SELU(), 1.
+    elif activation in ['softplus']:
+        return nn.Softplus(), 1
+    elif activation in ['softsign']:
+        return nn.Softsign(), 1  # unsure about this gain
+    elif activation in ['sigmoid', 'logistic']:
+        return nn.Sigmoid(), 1.
+    elif activation in ['tanh']:
+        return nn.Tanh(), 1.
+    else:
+        raise ValueError(
+            'Activation "{}" not available.'.format(activation)
+        )
+
+
+class Swish(nn.Module):
+    """
+    Performs the 'Swish' non-linear activation function.
+    https://arxiv.org/pdf/1710.05941.pdf
+    Arguments:
+        affine (bool): Multiply the input to sigmoid
+            with a learnable scale. Default value is False.
+    """
+    def __init__(self, affine=False):
+        super(Swish, self).__init__()
+        if affine:
+            self.beta = nn.Parameter(torch.tensor([1.]))
+        self.affine = affine
+
+    def forward(self, input, *args, **kwargs):
+        """
+        Apply the swish non-linear activation function
+        and return the results.
+        Arguments:
+            input (torch.Tensor)
+        Returns:
+            output (torch.Tensor)
+        """
+        x = input
+        if self.affine:
+            x *= self.beta
+        return x * torch.sigmoid(x)
+
+
+def _get_weight_and_coef(shape, lr_mul=1, weight_scale=True, gain=1, fill=None):
+    """
+    Get an intialized weight and its runtime coefficients as an nn.Parameter tensor.
+    Arguments:
+        shape (tuple, list): Shape of weight tensor.
+        lr_mul (float): The learning rate multiplier for
+            this weight. Default value is 1.
+        weight_scale (bool): Use weight scaling for equalized
+            learning rate. Default value is True.
+        gain (float): The gain of the weight. Default value is 1.
+        fill (float, optional): Instead of initializing the weight
+            with scaled normally distributed values, fill it with
+            this value. Useful for bias weights.
+    Returns:
+        weight (nn.Parameter)
+    """
+    fan_in = np.prod(shape[1:])
+    he_std = gain / np.sqrt(fan_in)
+
+    if weight_scale:
+        init_std = 1 / lr_mul
+        runtime_coef = he_std * lr_mul
+    else:
+        init_std = he_std / lr_mul
+        runtime_coef = lr_mul
+
+    weight = torch.empty(*shape)
+    if fill is None:
+        weight.normal_(0, init_std)
+    else:
+        weight.fill_(fill)
+    return nn.Parameter(weight), runtime_coef
+
+
+def _apply_conv(input, *args, transpose=False, **kwargs):
+    """
+    Perform a 1d, 2d or 3d convolution with specified
+    positional and keyword arguments. Which type of
+    convolution that is used depends on shape of data.
+    Arguments:
+        input (torch.Tensor): The input data for the
+            convolution.
+        *args: Positional arguments for the convolution.
+    Keyword Arguments:
+        transpose (bool): Transpose the convolution.
+            Default value is False
+        **kwargs: Keyword arguments for the convolution.
+    """
+    dim = input.dim() - 2
+    conv_fn = getattr(
+        F, 'conv{}{}d'.format('_transpose' if transpose else '', dim))
+    return conv_fn(input=input, *args, **kwargs)
+
+
+def _setup_mod_weight_for_t_conv(weight, in_channels, out_channels):
+    """
+    Reshape a modulated conv weight for use with a transposed convolution.
+    Arguments:
+        weight (torch.Tensor)
+        in_channels (int)
+        out_channels (int)
+    Returns:
+        reshaped_weight (torch.Tensor)
+    """
+    # [BO]I*k -> BOI*k
+    weight = weight.view(
+        -1,
+        out_channels,
+        in_channels,
+        *weight.size()[2:]
+    )
+    # BOI*k -> BIO*k
+    weight = weight.transpose(1, 2)
+    # BIO*k -> [BI]O*k
+    weight = weight.reshape(
+        -1,
+        out_channels,
+        *weight.size()[3:]
+    )
+    return weight
+
+
+def _setup_filter_kernel(filter_kernel, gain=1, up_factor=1, dim=2):
+    """
+    Set up a filter kernel and return it as a tensor.
+    Arguments:
+        filter_kernel (int, list, torch.tensor, None): The filter kernel
+            values to use. If this value is an int, a binomial filter of
+            this size is created. If a sequence with a single axis is used,
+            it will be expanded to the number of `dims` specified. If value
+            is None, a filter of values [1, 1] is used.
+        gain (float): Gain of the filter kernel. Default value is 1.
+        up_factor (int): Scale factor. Should only be given for upscaling filters.
+            Default value is 1.
+        dim (int): Number of dimensions of data. Default value is 2.
+    Returns:
+        filter_kernel_tensor (torch.Tensor)
+    """
+    filter_kernel = filter_kernel or 2
+    if isinstance(filter_kernel, (int, float)):
+        def binomial(n, k):
+            if k in [1, n]:
+                return 1
+            return np.math.factorial(n) / (np.math.factorial(k) * np.math.factorial(n - k))
+        filter_kernel = [binomial(filter_kernel, k) for k in range(1, filter_kernel + 1)]
+    if not torch.is_tensor(filter_kernel):
+        filter_kernel = torch.tensor(filter_kernel)
+    filter_kernel = filter_kernel.float()
+    if filter_kernel.dim() == 1:
+        _filter_kernel = filter_kernel.unsqueeze(0)
+        while filter_kernel.dim() < dim:
+            filter_kernel = torch.matmul(
+                filter_kernel.unsqueeze(-1), _filter_kernel)
+    assert all(filter_kernel.size(0) == s for s in filter_kernel.size())
+    filter_kernel /= filter_kernel.sum()
+    filter_kernel *= gain * up_factor ** 2
+    return filter_kernel.float()
+
+
+def _get_layer(layer_class, kwargs, wrap=False, noise=False):
+    """
+    Create a layer and wrap it in optional
+    noise and/or bias/activation layers.
+    Arguments:
+        layer_class: The class of the layer to construct.
+        kwargs (dict): The keyword arguments to use for constructing
+            the layer and optionally the bias/activaiton layer.
+        wrap (bool): Wrap the layer in an bias/activation layer and
+            optionally a noise injection layer. Default value is False.
+        noise (bool): Inject noise before the bias/activation wrapper.
+            This can only be done when `wrap=True`. Default value is False.
+    """
+    layer = layer_class(**kwargs)
+    if wrap:
+        if noise:
+            layer = NoiseInjectionWrapper(layer)
+        layer = BiasActivationWrapper(layer, **kwargs)
+    return layer
+
+
+class BiasActivationWrapper(nn.Module):
+    """
+    Wrap a module to add bias and non-linear activation
+    to any outputs of that module.
+    Arguments:
+        layer (nn.Module): The module to wrap.
+        features (int, optional): The number of features
+            of the output of the `layer`. This argument
+            has to be specified if `use_bias=True`.
+        use_bias (bool): Add bias to the output.
+            Default value is True.
+        activation (str, nn.Module, callable, optional):
+            non-linear activation function to use.
+            Unused if notspecified.
+        bias_init (float): Value to initialize bias
+            weight with. Default value is 0.
+        lr_mul (float): Learning rate multiplier of
+            the bias weight. Weights are scaled by
+            this value. Default value is 1.
+        weight_scale (float): Scale weights for
+            equalized learning rate.
+            Default value is True.
+    """
+    def __init__(self,
+                 layer,
+                 features=None,
+                 use_bias=True,
+                 activation='linear',
+                 bias_init=0,
+                 lr_mul=1,
+                 weight_scale=True,
+                 *args,
+                 **kwargs):
+        super(BiasActivationWrapper, self).__init__()
+        self.layer = layer
+        bias = None
+        bias_coef = None
+        if use_bias:
+            assert features, '`features` is required when using bias.'
+            bias, bias_coef = _get_weight_and_coef(
+                shape=[features],
+                lr_mul=lr_mul,
+                weight_scale=False,
+                fill=bias_init
+            )
+        self.register_parameter('bias', bias)
+        self.bias_coef = bias_coef
+        self.act, self.gain = get_activation(activation)
+
+    def forward(self, *args, **kwargs):
+        """
+        Forward all possitional and keyword arguments
+        to the layer wrapped by this module and add
+        bias (if set) and run through non-linear activation
+        function (if set).
+        Arguments:
+            *args (positional arguments)
+            **kwargs (keyword arguments)
+        Returns:
+            output (torch.Tensor)
+        """
+        x = self.layer(*args, **kwargs)
+        if self.bias is not None:
+            bias = self.bias.view(1, -1, *[1] * (x.dim() - 2))
+            if self.bias_coef != 1:
+                bias = self.bias_coef * bias
+            x += bias
+        x = self.act(x)
+        if self.gain != 1:
+            x *= self.gain
+        return x
+
+    def extra_repr(self):
+        return 'bias={}'.format(self.bias is not None)
+
+
+class NoiseInjectionWrapper(nn.Module):
+    """
+    Wrap a module to add noise scaled by a
+    learnable parameter to any outputs of the
+    wrapped module.
+    Noise is randomized for each output but can
+    be set to static noise by calling `static_noise()`
+    of this object. This can only be done once data
+    has passed through this layer at least once so that
+    the shape of the static noise to create is known.
+    Check if the shape is known by calling `has_noise_shape()`.
+    Arguments:
+        layer (nn.Module): The module to wrap.
+        same_over_batch (bool): Repeat the same
+            noise values over the entire batch
+            instead of creating separate noise
+            values for each entry in the batch.
+            Default value is True.
+    """
+
+    def __init__(self, layer, same_over_batch=True):
+        super(NoiseInjectionWrapper, self).__init__()
+        self.layer = layer
+        self.weight = torch.nn.Parameter(torch.zeros(1))
+        self.register_buffer('noise_storage', None)
+        self.same_over_batch = same_over_batch
+        self.random_noise()
+
+    def has_noise_shape(self):
+        """
+        If this module has had data passed through it
+        the noise shape is known and this function returns
+        True. Else False.
+        Returns:
+            noise_shape_known (bool)
+        """
+        return self.noise_storage is not None
+
+    def random_noise(self):
+        """
+        Randomize noise for each
+        new output.
+        """
+        self._fixed_noise = False
+        if isinstance(self.noise_storage, nn.Parameter):
+            noise_storage = self.noise_storage
+            del self.noise_storage
+            self.register_buffer('noise_storage', noise_storage.data)
+
+    def static_noise(self, trainable=False, noise_tensor=None):
+        """
+        Set up static noise that can optionally be a trainable
+        parameter. Static noise does not change between inputs
+        unless the user has altered its values. Returns the tensor
+        object that stores the static noise.
+        Arguments:
+            trainable (bool): Wrap the static noise tensor in
+                nn.Parameter to make it trainable. The returned
+                tensor will be wrapped.
+            noise_tensor (torch.Tensor, optional): A predefined
+                static noise tensor. If not passed, one will be
+                created.
+        """
+        assert self.has_noise_shape(), \
+            'Noise shape is unknown'
+        if noise_tensor is None:
+            noise_tensor = self.noise_storage
+        else:
+            noise_tensor = noise_tensor.to(self.weight)
+        if trainable and not isinstance(noise_tensor, nn.Parameter):
+            noise_tensor = nn.Parameter(noise_tensor)
+        if isinstance(self.noise_storage, nn.Parameter) and not trainable:
+            del self.noise_storage
+            self.register_buffer('noise_storage', noise_tensor)
+        else:
+            self.noise_storage = noise_tensor
+        self._fixed_noise = True
+        return noise_tensor
+
+    def _save_to_state_dict(self, destination, prefix, keep_vars):
+        r"""Saves module state to `destination` dictionary, containing a state
+        submodule in :meth:`~torch.nn.Module.state_dict`.
+
+        Overridden to ignore the noise storage buffer.
+
+        Arguments:
+            destination (dict): a dict where state will be stored
+            prefix (str): the prefix for parameters and buffers used in this
+                module
+        """
+        for name, param in self._parameters.items():
+            if name != 'noise_storage' and param is not None:
+                destination[prefix + name] = param if keep_vars else param.data
+        for name, buf in self._buffers.items():
+            if name != 'noise_storage' and buf is not None:
+                destination[prefix + name] = buf if keep_vars else buf.data
+
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        r"""Copies parameters and buffers from :attr:`state_dict` into only
+        this module, but not its descendants. This is called on every submodule
+        in :meth:`~torch.nn.Module.load_state_dict`. Metadata saved for this
+        module in input :attr:`state_dict` is provided as :attr:`local_metadata`.
+        For state dicts without metadata, :attr:`local_metadata` is empty.
+        Overridden to ignore noise storage buffer.
+        """
+        key = prefix + 'noise_storage'
+        if key in state_dict:
+            del state_dict[key]
+        return super(NoiseInjectionWrapper, self)._load_from_state_dict(
+            state_dict, prefix, *args, **kwargs)
+
+    def forward(self, *args, **kwargs):
+        """
+        Forward all possitional and keyword arguments
+        to the layer wrapped by this module and add
+        noise to its outputs before returning them.
+        Arguments:
+            *args (positional arguments)
+            **kwargs (keyword arguments)
+        Returns:
+            output (torch.Tensor)
+        """
+        x = self.layer(*args, **kwargs)
+        noise_shape = list(x.size())
+        noise_shape[1] = 1
+        if self.same_over_batch:
+            noise_shape[0] = 1
+        if self.noise_storage is None or list(self.noise_storage.size()) != noise_shape:
+            if not self._fixed_noise:
+                self.noise_storage = torch.empty(
+                    *noise_shape,
+                    dtype=self.weight.dtype,
+                    device=self.weight.device
+                )
+            else:
+                assert list(self.noise_storage.size()[2:]) == noise_shape[2:], \
+                    'A data size {} has been encountered, '.format(x.size()[2:]) + \
+                    'the static noise previously set up does ' + \
+                    'not match this size {}'.format(self.noise_storage.size()[2:])
+                assert self.noise_storage.size(0) == 1 or self.noise_storage.size(0) == x.size(0), \
+                    'Static noise batch size mismatch! ' + \
+                    'Noise batch size: {}, '.format(self.noise_storage.size(0)) + \
+                    'input batch size: {}'.format(x.size(0))
+                assert self.noise_storage.size(1) == 1 or self.noise_storage.size(1) == x.size(1), \
+                    'Static noise channel size mismatch! ' + \
+                    'Noise channel size: {}, '.format(self.noise_storage.size(1)) + \
+                    'input channel size: {}'.format(x.size(1))
+        if not self._fixed_noise:
+            self.noise_storage.normal_()
+        x += self.weight * self.noise_storage
+        return x
+
+    def extra_repr(self):
+        return 'static_noise={}'.format(self._fixed_noise)
+
+
+class FilterLayer(nn.Module):
+    """
+    Apply a filter by using convolution.
+    Arguments:
+        filter_kernel (torch.Tensor): The filter kernel to use.
+            Should be of shape `dims * (k,)` where `k` is the
+            kernel size and `dims` is the number of data dimensions
+            (excluding batch and channel dimension).
+        stride (int): The stride of the convolution.
+        pad0 (int): Amount to pad start of each data dimension.
+            Default value is 0.
+        pad1 (int): Amount to pad end of each data dimension.
+            Default value is 0.
+        pad_mode (str): The padding mode. Default value is 'constant'.
+        pad_constant (float): The constant value to pad with if
+            `pad_mode='constant'`. Default value is 0.
+    """
+    def __init__(self,
+                 filter_kernel,
+                 stride=1,
+                 pad0=0,
+                 pad1=0,
+                 pad_mode='constant',
+                 pad_constant=0,
+                 *args,
+                 **kwargs):
+        super(FilterLayer, self).__init__()
+        dim = filter_kernel.dim()
+        filter_kernel = filter_kernel.view(1, 1, *filter_kernel.size())
+        self.register_buffer('filter_kernel', filter_kernel)
+        self.stride = stride
+        if pad0 == pad1 and (pad0 == 0 or pad_mode == 'constant' and pad_constant == 0):
+            self.fused_pad = True
+            self.padding = pad0
+        else:
+            self.fused_pad = False
+            self.padding = [pad0, pad1] * dim
+            self.pad_mode = pad_mode
+            self.pad_constant = pad_constant
+
+    def forward(self, input, **kwargs):
+        """
+        Pad the input and run the filter over it
+        before returning the new values.
+        Arguments:
+            input (torch.Tensor)
+        Returns:
+            output (torch.Tensor)
+        """
+        x = input
+        conv_kwargs = dict(
+            weight=self.filter_kernel.repeat(
+                input.size(1), *[1] * (self.filter_kernel.dim() - 1)),
+            stride=self.stride,
+            groups=input.size(1),
+        )
+        if self.fused_pad:
+            conv_kwargs.update(padding=self.padding)
+        else:
+            x = F.pad(x, self.padding, mode=self.pad_mode, value=self.pad_constant)
+        return _apply_conv(
+            input=x,
+            transpose=False,
+            **conv_kwargs
+        )
+
+    def extra_repr(self):
+        return 'filter_size={}, stride={}'.format(
+            tuple(self.filter_kernel.size()[2:]), self.stride)
+
+
+class Upsample(nn.Module):
+    """
+    Performs upsampling without learnable parameters that doubles
+    the size of data.
+    Arguments:
+        mode (str): 'FIR' or one of the valid modes
+            that can be passed to torch.nn.functional.interpolate().
+        filter (int, list, tensor): Filter to use if `mode='FIR'`.
+            Default value is a lowpass filter of values [1, 3, 3, 1].
+        filter_pad_mode (str): If `mode='FIR'`, this is used with the filter.
+            See `FilterLayer` docstring for more info.
+        filter_pad_constant (float): If `mode='FIR'`, this is used with the filter.
+            See `FilterLayer` docstring for more info.
+        gain (float): If `mode='FIR'`, this is used with the filter.
+            See `FilterLayer` docstring for more info.
+        dim (int): Dims of data (excluding batch and channel dimensions).
+            Default value is 2.
+    """
+
+    def __init__(self,
+                 mode='FIR',
+                 filter=[1, 3, 3, 1],
+                 filter_pad_mode='constant',
+                 filter_pad_constant=0,
+                 gain=1,
+                 dim=2,
+                 *args,
+                 **kwargs):
+        super(Upsample, self).__init__()
+        assert mode != 'max', 'mode \'max\' can only be used for downsampling.'
+        if mode == 'FIR':
+            if filter is None:
+                filter = [1, 1]
+            filter_kernel = _setup_filter_kernel(
+                filter_kernel=filter,
+                gain=gain,
+                up_factor=2,
+                dim=dim
+            )
+            pad = filter_kernel.size(-1) - 1
+            self.filter = FilterLayer(
+                filter_kernel=filter_kernel,
+                pad0=(pad + 1) // 2 + 1,
+                pad1=pad // 2,
+                pad_mode=filter_pad_mode,
+                pad_constant=filter_pad_constant
+            )
+            self.register_buffer('weight', torch.ones(*[1 for _ in range(dim + 2)]))
+        self.mode = mode
+
+    def forward(self, input, **kwargs):
+        """
+        Upsample inputs.
+        Arguments:
+            input (torch.Tensor)
+        Returns:
+            output (torch.Tensor)
+        """
+        if self.mode == 'FIR':
+            x = _apply_conv(
+                input=input,
+                weight=self.weight.expand(input.size(1), *self.weight.size()[1:]),
+                groups=input.size(1),
+                stride=2,
+                transpose=True
+            )
+            x = self.filter(x)
+        else:
+            interp_kwargs = dict(scale_factor=2, mode=self.mode)
+            if 'linear' in self.mode or 'cubic' in self.mode:
+                interp_kwargs.update(align_corners=False)
+            x = F.interpolate(input, **interp_kwargs)
+        return x
+
+    def extra_repr(self):
+        return 'resample_mode={}'.format(self.mode)
+
+
+class Downsample(nn.Module):
+    """
+    Performs downsampling without learnable parameters that
+    reduces size of data by half.
+    Arguments:
+        mode (str): 'FIR', 'max' or one of the valid modes
+            that can be passed to torch.nn.functional.interpolate().
+        filter (int, list, tensor): Filter to use if `mode='FIR'`.
+            Default value is a lowpass filter of values [1, 3, 3, 1].
+        filter_pad_mode (str): If `mode='FIR'`, this is used with the filter.
+            See `FilterLayer` docstring for more info.
+        filter_pad_constant (float): If `mode='FIR'`, this is used with the filter.
+            See `FilterLayer` docstring for more info.
+        gain (float): If `mode='FIR'`, this is used with the filter.
+            See `FilterLayer` docstring for more info.
+        dim (int): Dims of data (excluding batch and channel dimensions).
+            Default value is 2.
+    """
+
+    def __init__(self,
+                 mode='FIR',
+                 filter=[1, 3, 3, 1],
+                 filter_pad_mode='constant',
+                 filter_pad_constant=0,
+                 gain=1,
+                 dim=2,
+                 *args,
+                 **kwargs):
+        super(Downsample, self).__init__()
+        if mode == 'FIR':
+            if filter is None:
+                filter = [1, 1]
+            filter_kernel = _setup_filter_kernel(
+                filter_kernel=filter,
+                gain=gain,
+                up_factor=1,
+                dim=dim
+            )
+            pad = filter_kernel.size(-1) - 2
+            pad0 = pad // 2
+            pad1 = pad - pad0
+            self.filter = FilterLayer(
+                filter_kernel=filter_kernel,
+                stride=2,
+                pad0=pad0,
+                pad1=pad1,
+                pad_mode=filter_pad_mode,
+                pad_constant=filter_pad_constant
+            )
+        self.mode = mode
+
+    def forward(self, input, **kwargs):
+        """
+        Downsample inputs to half its size.
+        Arguments:
+            input (torch.Tensor)
+        Returns:
+            output (torch.Tensor)
+        """
+        if self.mode == 'FIR':
+            x = self.filter(input)
+        elif self.mode == 'max':
+            return getattr(F, 'max_pool{}d'.format(input.dim() - 2))(input)
+        else:
+            x = F.interpolate(input, scale_factor=0.5, mode=self.mode)
+        return x
+
+    def extra_repr(self):
+        return 'resample_mode={}'.format(self.mode)
+
+
+class MinibatchStd(nn.Module):
+    """
+    Adds the aveage std of each data point over a
+    slice of the minibatch to that slice as a new
+    feature map. This gives an output with one extra
+    channel.
+    Arguments:
+        group_size (int): Number of entries in each slice
+            of the batch. If <= 0, the entire batch is used.
+            Default value is 4.
+        eps (float): Epsilon value added for numerical stability.
+            Default value is 1e-8.
+    """
+    def __init__(self, group_size=4, eps=1e-8, *args, **kwargs):
+        super(MinibatchStd, self).__init__()
+        if group_size is None or group_size <= 0:
+            # Entire batch as group size
+            group_size = 0
+        assert group_size != 1, 'Can not use 1 as minibatch std group size.'
+        self.group_size = group_size
+        self.eps = eps
+
+    def forward(self, input, **kwargs):
+        """
+        Add a new feature map to the input containing the average
+        standard deviation for each slice.
+        Arguments:
+            input (torch.Tensor)
+        Returns:
+            output (torch.Tensor)
+        """
+        group_size = self.group_size or input.size(0)
+        assert input.size(0) >= group_size, \
+            'Can not use a smaller batch size ' + \
+            '({}) than the specified '.format(input.size(0)) + \
+            'group size ({}) '.format(group_size) + \
+            'of this minibatch std layer.'
+        assert input.size(0) % group_size == 0, \
+            'Can not use a batch of a size ' + \
+            '({}) that is not '.format(input.size(0)) + \
+            'evenly divisible by the group size ({})'.format(group_size)
+        x = input
+
+        # B = batch size, C = num channels
+        # *s = the size dimensions (height, width for images)
+
+        # BC*s -> G[B/G]C*s
+        y = input.view(group_size, -1, *input.size()[1:])
+        # For numerical stability when training with mixed precision
+        y = y.float()
+        # G[B/G]C*s
+        y -= y.mean(dim=0, keepdim=True)
+        # [B/G]C*s
+        y = torch.mean(y ** 2, dim=0)
+        # [B/G]C*s
+        y = torch.sqrt(y + self.eps)
+        # [B/G]
+        y = torch.mean(y.view(y.size(0), -1), dim=-1)
+        # [B/G]1*1
+        y = y.view(-1, *[1] * (input.dim() - 1))
+        # Cast back to input dtype
+        y = y.to(x)
+        # B1*1
+        y = y.repeat(group_size, *[1] * (y.dim() - 1))
+        # B1*s
+        y = y.expand(y.size(0), 1, *x.size()[2:])
+        # B[C+1]*s
+        x = torch.cat([x, y], dim=1)
+        return x
+
+    def extra_repr(self):
+        return 'group_size={}'.format(self.group_size or '-1')
+
+
+class DenseLayer(nn.Module):
+    """
+    A fully connected layer.
+    NOTE: No bias is applied in this layer.
+    Arguments:
+        in_features (int): Number of input features.
+        out_features (int): Number of output features.
+        lr_mul (float): Learning rate multiplier of
+            the weight. Weights are scaled by
+            this value. Default value is 1.
+        weight_scale (float): Scale weights for
+            equalized learning rate.
+            Default value is True.
+        gain (float): The gain of this layer. Default value is 1.
+    """
+    def __init__(self,
+                 in_features,
+                 out_features,
+                 lr_mul=1,
+                 weight_scale=True,
+                 gain=1,
+                 *args,
+                 **kwargs):
+        super(DenseLayer, self).__init__()
+        weight, weight_coef = _get_weight_and_coef(
+            shape=[out_features, in_features],
+            lr_mul=lr_mul,
+            weight_scale=weight_scale,
+            gain=gain
+        )
+        self.register_parameter('weight', weight)
+        self.weight_coef = weight_coef
+
+    def forward(self, input, **kwargs):
+        """
+        Perform a matrix multiplication of the weight
+        of this layer and the input.
+        Arguments:
+            input (torch.Tensor)
+        Returns:
+            output (torch.Tensor)
+        """
+        weight = self.weight
+        if self.weight_coef != 1:
+            weight = self.weight_coef * weight
+        return input.matmul(weight.t())
+
+    def extra_repr(self):
+        return 'in_features={}, out_features={}'.format(
+            self.weight.size(1), self.weight.size(0))
+
+
+class ConvLayer(nn.Module):
+    """
+    A convolutional layer that can have its outputs
+    modulated (style mod). It can also normalize outputs.
+    These operations are done by modifying the convolutional
+    kernel weight and employing grouped convolutions for
+    efficiency.
+    NOTE: No bias is applied in this layer.
+    NOTE: Amount of padding used is the same as 'SAME'
+        argument in tensorflow for conv padding.
+    Arguments:
+        in_channels (int): Number of input channels.
+        out_channels (int): Number of output channels.
+        latent_size (int, optional): The size of the
+            latents to use for modulating this convolution.
+            Only required when `modulate=True`.
+        modulate (bool): Applies a "style" to the outputs
+            of the layer. The style is given by a latent
+            vector passed with the input to this layer.
+            A dense layer is added that projects the
+            values of the latent into scales for the
+            data channels.
+            Default value is False.
+        demodulate (bool): Normalize std of outputs.
+            Can only be set to True when `modulate=True`.
+            Default value is False.
+        kernel_size (int): The size of the kernel.
+            Default value is 3.
+        pad_mode (str): The padding mode. Default value is 'constant'.
+        pad_constant (float): The constant value to pad with if
+            `pad_mode='constant'`. Default value is 0.
+        lr_mul (float): Learning rate multiplier of
+            the weight. Weights are scaled by
+            this value. Default value is 1.
+        weight_scale (float): Scale weights for
+            equalized learning rate.
+            Default value is True.
+        gain (float): The gain of this layer. Default value is 1.
+        dim (int): Dims of data (excluding batch and channel dimensions).
+            Default value is 2.
+        eps (float): Epsilon value added for numerical stability.
+            Default value is 1e-8.
+    """
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 latent_size=None,
+                 modulate=False,
+                 demodulate=False,
+                 kernel_size=3,
+                 pad_mode='constant',
+                 pad_constant=0,
+                 lr_mul=1,
+                 weight_scale=True,
+                 gain=1,
+                 dim=2,
+                 eps=1e-8,
+                 *args,
+                 **kwargs):
+        super(ConvLayer, self).__init__()
+        assert modulate or not demodulate, '`demodulate=True` can ' + \
+            'only be used when `modulate=True`'
+        if modulate:
+            assert latent_size is not None, 'When using `modulate=True`, ' + \
+                '`latent_size` has to be specified.'
+        kernel_shape = [out_channels, in_channels] + dim * [kernel_size]
+        weight, weight_coef = _get_weight_and_coef(
+            shape=kernel_shape,
+            lr_mul=lr_mul,
+            weight_scale=weight_scale,
+            gain=gain
+        )
+        self.register_parameter('weight', weight)
+        self.weight_coef = weight_coef
+        if modulate:
+            self.dense = BiasActivationWrapper(
+                layer=DenseLayer(
+                    in_features=latent_size,
+                    out_features=in_channels,
+                    lr_mul=lr_mul,
+                    weight_scale=weight_scale,
+                    gain=1
+                ),
+                features=in_channels,
+                use_bias=True,
+                activation='linear',
+                bias_init=1,
+                lr_mul=lr_mul,
+                weight_scale=weight_scale,
+            )
+        self.dense_reshape = [-1, 1, in_channels] + dim * [1]
+        self.dmod_reshape = [-1, out_channels, 1] + dim * [1]
+        pad = (kernel_size - 1)
+        pad0 = pad - pad // 2
+        pad1 = pad - pad0
+        if pad0 == pad1 and (pad0 == 0 or pad_mode == 'constant' and pad_constant == 0):
+            self.fused_pad = True
+            self.padding = pad0
+        else:
+            self.fused_pad = False
+            self.padding = [pad0, pad1] * dim
+        self.pad_mode = pad_mode
+        self.pad_constant = pad_constant
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.latent_size = latent_size
+        self.modulate = modulate
+        self.demodulate = demodulate
+        self.kernel_size = kernel_size
+        self.lr_mul = lr_mul
+        self.weight_scale = weight_scale
+        self.gain = gain
+        self.dim = dim
+        self.eps = eps
+
+    def forward_mod(self, input, latent, weight, **kwargs):
+        """
+        Run the forward operation with modulation.
+        Automatically called from `forward()` if modulation
+        is enabled.
+        """
+        assert latent is not None, 'A latent vector is ' + \
+            'required for the forwad pass of a modulated conv layer.'
+
+        # B = batch size, C = num channels
+        # *s = the size dimensions, example: (height, width) for images
+        # *k = sizes of the convolutional kernel excluding in and out channel dimensions.
+        # *1 = multiple dimensions of size 1, with number of dimensions depending on data format.
+        # O = num output channels, I = num input channels
+
+        # BI
+        style_mod = self.dense(input=latent)
+        # B1I*1
+        style_mod = style_mod.view(*self.dense_reshape)
+        # 1OI*k
+        weight = weight.unsqueeze(0)
+        # (1OI*k)x(B1I*1) -> BOI*k
+        weight = weight * style_mod
+        if self.demodulate:
+            # BO
+            dmod = torch.rsqrt(
+                torch.sum(
+                    weight.view(
+                        weight.size(0),
+                        weight.size(1),
+                        -1
+                    ) ** 2,
+                    dim=-1
+                ) + self.eps
+            )
+            # BO1*1
+            dmod = dmod.view(*self.dmod_reshape)
+            # (BOI*k)x(BO1*1) -> BOI*k
+            weight = weight * dmod
+        # BI*s -> 1[BI]*s
+        x = input.view(1, -1, *input.size()[2:])
+        # BOI*k -> [BO]I*k
+        weight = weight.view(-1, *weight.size()[2:])
+        # 1[BO]*s
+        x = self._process(input=x, weight=weight, groups=input.size(0))
+        # 1[BO]*s -> BO*s
+        x = x.view(-1, self.out_channels, *x.size()[2:])
+        return x
+
+    def forward(self, input, latent=None, **kwargs):
+        """
+        Convolve the input.
+        Arguments:
+            input (torch.Tensor)
+            latents (torch.Tensor, optional)
+        Returns:
+            output (torch.Tensor)
+        """
+        weight = self.weight
+        if self.weight_coef != 1:
+            weight = self.weight_coef * weight
+        if self.modulate:
+            return self.forward_mod(input=input, latent=latent, weight=weight)
+        return self._process(input=input, weight=weight)
+
+    def _process(self, input, weight, **kwargs):
+        """
+        Pad input and convolve it returning the result.
+        """
+        x = input
+        if self.fused_pad:
+            kwargs.update(padding=self.padding)
+        else:
+            x = F.pad(x, self.padding, mode=self.pad_mode, value=self.pad_constant)
+        return _apply_conv(input=x, weight=weight, transpose=False, **kwargs)
+
+    def extra_repr(self):
+        string = 'in_channels={}, out_channels={}'.format(
+            self.weight.size(1), self.weight.size(0))
+        string += ', modulate={}, demodulate={}'.format(
+            self.modulate, self.demodulate)
+        return string
+
+
+class ConvUpLayer(ConvLayer):
+    """
+    A convolutional upsampling layer that doubles the size of inputs.
+    Extends the functionality of the `ConvLayer` class.
+    Arguments:
+        Same arguments as the `ConvLayer` class.
+    Class Specific Keyword Arguments:
+        fused (bool): Fuse the upsampling operation with the
+            convolution, turning this layer into a strided transposed
+            convolution. Default value is True.
+        mode (str): Resample mode, can only be 'FIR' or 'none' if the operation
+            is fused, otherwise it can also be one of the valid modes
+            that can be passed to torch.nn.functional.interpolate().
+        filter (int, list, tensor): Filter to use if `mode='FIR'`.
+            Default value is a lowpass filter of values [1, 3, 3, 1].
+        filter_pad_mode (str): If `mode='FIR'`, this is used with the filter.
+            See `FilterLayer` docstring for more info.
+        filter_pad_constant (float): If `mode='FIR'`, this is used with the filter.
+            See `FilterLayer` docstring for more info.
+        pad_once (bool): If FIR filter is used, do all the padding for
+            both convolution and FIR in the FIR layer instead of once per layer.
+            Default value is True.
+    """
+
+    def __init__(self,
+                 *args,
+                 fused=True,
+                 mode='FIR',
+                 filter=[1, 3, 3, 1],
+                 filter_pad_mode='constant',
+                 filter_pad_constant=0,
+                 pad_once=True,
+                 **kwargs):
+        super(ConvUpLayer, self).__init__(*args, **kwargs)
+        if fused:
+            assert mode in ['FIR', 'none'], \
+                'Fused conv upsample can only use ' + \
+                '\'FIR\' or \'none\' for resampling ' + \
+                '(`mode` argument).'
+            self.padding = np.ceil(self.kernel_size / 2 - 1)
+            self.output_padding = 2 * (self.padding + 1) - self.kernel_size
+            if not self.modulate:
+                # pre-prepare weights only once instead of every forward pass
+                self.weight = nn.Parameter(self.weight.data.transpose(0, 1).contiguous())
+            self.filter = None
+            if mode == 'FIR':
+                filter_kernel = _setup_filter_kernel(
+                    filter_kernel=filter,
+                    gain=self.gain,
+                    up_factor=2,
+                    dim=self.dim
+                )
+                if pad_once:
+                    self.padding = 0
+                    self.output_padding = 0
+                    pad = (filter_kernel.size(-1) - 2) - (self.kernel_size - 1)
+                    pad0 = (pad + 1) // 2 + 1,
+                    pad1 = pad // 2 + 1,
+                else:
+                    pad = (filter_kernel.size(-1) - 1)
+                    pad0 = pad // 2
+                    pad1 = pad - pad0
+                self.filter = FilterLayer(
+                    filter_kernel=filter_kernel,
+                    pad0=pad0,
+                    pad1=pad1,
+                    pad_mode=filter_pad_mode,
+                    pad_constant=filter_pad_constant
+                )
+        else:
+            assert mode != 'none', '\'none\' can not be used as ' + \
+                'sampling `mode` when `fused=False` as upsampling ' + \
+                'has to be performed separately from the conv layer.'
+            self.upsample = Upsample(
+                mode=mode,
+                filter=filter,
+                filter_pad_mode=filter_pad_mode,
+                filter_pad_constant=filter_pad_constant,
+                channels=self.in_channels,
+                gain=self.gain,
+                dim=self.dim
+            )
+        self.fused = fused
+        self.mode = mode
+
+    def _process(self, input, weight, **kwargs):
+        """
+        Apply resampling (if enabled) and convolution.
+        """
+        x = input
+        if self.fused:
+            if self.modulate:
+                weight = _setup_mod_weight_for_t_conv(
+                    weight=weight,
+                    in_channels=self.in_channels,
+                    out_channels=self.out_channels
+                )
+            pad_out = False
+            if self.pad_mode == 'constant' and self.pad_constant == 0:
+                if self.filter is not None or not self.pad_once:
+                    kwargs.update(
+                        padding=self.padding,
+                        output_padding=self.output_padding,
+                    )
+            elif self.filter is None:
+                if self.padding:
+                    x = F.pad(
+                        x,
+                        [self.padding] * 2 * self.dim,
+                        mode=self.pad_mode,
+                        value=self.pad_constant
+                    )
+                pad_out = self.output_padding != 0
+            kwargs.update(stride=2)
+            x = _apply_conv(
+                input=x,
+                weight=weight,
+                transpose=True,
+                **kwargs
+            )
+            if pad_out:
+                x = F.pad(
+                    x,
+                    [self.output_padding, 0] * self.dim,
+                    mode=self.pad_mode,
+                    value=self.pad_constant
+                )
+            if self.filter is not None:
+                x = self.filter(x)
+        else:
+            x = super(ConvUpLayer, self)._process(
+                input=self.upsample(input=x),
+                weight=weight,
+                **kwargs
+            )
+        return x
+
+    def extra_repr(self):
+        string = super(ConvUpLayer, self).extra_repr()
+        string += ', fused={}, resample_mode={}'.format(
+            self.fused, self.mode)
+        return string
+
+
+class ConvDownLayer(ConvLayer):
+    """
+    A convolutional downsampling layer that halves the size of inputs.
+    Extends the functionality of the `ConvLayer` class.
+    Arguments:
+        Same arguments as the `ConvLayer` class.
+    Class Specific Keyword Arguments:
+        fused (bool): Fuse the downsampling operation with the
+            convolution, turning this layer into a strided convolution.
+            Default value is True.
+        mode (str): Resample mode, can only be 'FIR' or 'none' if the operation
+            is fused, otherwise it can also be 'max' or one of the valid modes
+            that can be passed to torch.nn.functional.interpolate().
+        filter (int, list, tensor): Filter to use if `mode='FIR'`.
+            Default value is a lowpass filter of values [1, 3, 3, 1].
+        filter_pad_mode (str): If `mode='FIR'`, this is used with the filter.
+            See `FilterLayer` docstring for more info.
+        filter_pad_constant (float): If `mode='FIR'`, this is used with the filter.
+            See `FilterLayer` docstring for more info.
+        pad_once (bool): If FIR filter is used, do all the padding for
+            both convolution and FIR in the FIR layer instead of once per layer.
+            Default value is True.
+    """
+
+    def __init__(self,
+                 *args,
+                 fused=True,
+                 mode='FIR',
+                 filter=[1, 3, 3, 1],
+                 filter_pad_mode='constant',
+                 filter_pad_constant=0,
+                 pad_once=True,
+                 **kwargs):
+        super(ConvDownLayer, self).__init__(*args, **kwargs)
+        if fused:
+            assert mode in ['FIR', 'none'], \
+                'Fused conv downsample can only use ' + \
+                '\'FIR\' or \'none\' for resampling ' + \
+                '(`mode` argument).'
+            pad = self.kernel_size - 2
+            pad0 = pad // 2
+            pad1 = pad - pad0
+            if pad0 == pad1 and (pad0 == 0 or self.pad_mode == 'constant' and self.pad_constant == 0):
+                self.fused_pad = True
+                self.padding = pad0
+            else:
+                self.fused_pad = False
+                self.padding = [pad0, pad1] * self.dim
+            self.filter = None
+            if mode == 'FIR':
+                filter_kernel = _setup_filter_kernel(
+                    filter_kernel=filter,
+                    gain=self.gain,
+                    up_factor=1,
+                    dim=self.dim
+                )
+                if pad_once:
+                    self.fused_pad = True
+                    self.padding = 0
+                    pad = (filter_kernel.size(-1) - 2) + (self.kernel_size - 1)
+                    pad0 = (pad + 1) // 2,
+                    pad1 = pad // 2,
+                else:
+                    pad = (filter_kernel.size(-1) - 1)
+                    pad0 = pad // 2
+                    pad1 = pad - pad0
+                self.filter = FilterLayer(
+                    filter_kernel=filter_kernel,
+                    pad0=pad0,
+                    pad1=pad1,
+                    pad_mode=filter_pad_mode,
+                    pad_constant=filter_pad_constant
+                )
+                self.pad_once = pad_once
+        else:
+            assert mode != 'none', '\'none\' can not be used as ' + \
+                'sampling `mode` when `fused=False` as downsampling ' + \
+                'has to be performed separately from the conv layer.'
+            self.downsample = Downsample(
+                mode=mode,
+                filter=filter,
+                pad_mode=filter_pad_mode,
+                pad_constant=filter_pad_constant,
+                channels=self.in_channels,
+                gain=self.gain,
+                dim=self.dim
+            )
+        self.fused = fused
+        self.mode = mode
+
+    def _process(self, input, weight, **kwargs):
+        """
+        Apply resampling (if enabled) and convolution.
+        """
+        x = input
+        if self.fused:
+            kwargs.update(stride=2)
+            if self.filter is not None:
+                x = self.filter(input=x)
+        else:
+            x = self.downsample(input=x)
+        x = super(ConvDownLayer, self)._process(
+            input=x,
+            weight=weight,
+            **kwargs
+        )
+        return x
+
+    def extra_repr(self):
+        string = super(ConvDownLayer, self).extra_repr()
+        string += ', fused={}, resample_mode={}'.format(
+            self.fused, self.mode)
+        return string
+
+
+class GeneratorConvBlock(nn.Module):
+    """
+    A convblock for the synthesiser model.
+    Arguments:
+        in_channels (int): Number of input channels.
+        out_channels (int): Number of output channels.
+        latent_size (int): The size of the latent vectors.
+        demodulate (bool): Normalize feature outputs from conv
+            layers. Default value is True.
+        resnet (bool): Use residual connections. Default value is
+            False.
+        up (bool): Upsample the data to twice its size. This is
+            performed in the first layer of the block. Default
+            value is False.
+        num_layers (int): Number of convolutional layers of this
+            block. Default value is 2.
+        filter (int, list): The filter to use if
+            `up=True` and `mode='FIR'`. If int, a low
+            pass filter of this size will be used. If list,
+            the filter is explicitly specified. If the filter
+            is of a single dimension it will be expanded to
+            the number of dimensions of the data. Default
+            value is a low pass filter of [1, 3, 3, 1].
+        activation (str, callable, nn.Module): The non-linear
+            activation function to use.
+            Default value is leaky relu with a slope of 0.2.
+        mode (str): The resample mode of upsampling layers.
+            Only used when `up=True`. If fused=True` only 'FIR'
+            and 'none' can be used. Else, anything that can
+            be passed to torch.nn.functional.interpolate is
+            a valid mode. Default value is 'FIR'.
+        fused (bool): If `up=True`, fuse the upsample operation
+            and the first convolutional layer into a transposed
+            convolutional layer.
+        kernel_size (int): Size of the convolutional kernel.
+            Default value is 3.
+        pad_mode (str): The padding mode for convolutional
+            layers. Has to be one of 'constant', 'reflect',
+            'replicate' or 'circular'. Default value is
+            'constant'.
+        pad_constant (float): The value to use for conv
+            padding if `conv_pad_mode='constant'`. Default
+            value is 0.
+        filter_pad_mode (str): If `mode='FIR'`, this is used with the filter.
+            Otherwise works the same as `pad_mode`.
+        filter_pad_constant (float): If `mode='FIR'`, this is used with the filter.
+            Otherwise works the same as `pad_constant`
+        pad_once (bool): If FIR filter is used, do all the padding for
+            both convolution and FIR in the FIR layer instead of once per layer.
+            Default value is True.
+        use_bias (bool): Add bias to layer outputs. Default value is True.
+        noise (bool): Add noise to the output of each layer. Default value
+            is True.
+        lr_mul (float): The learning rate multiplier for this
+            block. When loading weights of previously trained
+            networks, this value has to be the same as when
+            the network was trained for the outputs to not
+            change (as this is used to scale the weights).
+            Default value is 1.
+        weight_scale (bool): Use weight scaling for
+            equalized learning rate. Default value
+            is True.
+        eps (float): Epsilon value added for numerical stability.
+            Default value is 1e-8.
+    """
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 latent_size,
+                 demodulate=True,
+                 resnet=False,
+                 up=False,
+                 num_layers=2,
+                 filter=[1, 3, 3, 1],
+                 activation='leaky:0.2',
+                 mode='FIR',
+                 fused=True,
+                 kernel_size=3,
+                 pad_mode='constant',
+                 pad_constant=0,
+                 filter_pad_mode='constant',
+                 filter_pad_constant=0,
+                 pad_once=True,
+                 use_bias=True,
+                 noise=True,
+                 lr_mul=1,
+                 weight_scale=True,
+                 gain=1,
+                 dim=2,
+                 eps=1e-8,
+                 *args,
+                 **kwargs):
+        super(GeneratorConvBlock, self).__init__()
+        layer_kwargs = locals()
+        layer_kwargs.pop('self')
+        layer_kwargs.pop('__class__')
+        layer_kwargs.update(
+            features=out_channels,
+            modulate=True,
+        )
+
+        assert num_layers > 0
+        assert 1 <= dim <= 3, '`dim` can only be 1, 2 or 3.'
+        if up:
+            available_sampling = ['FIR']
+            if fused:
+                available_sampling.append('none')
+            else:
+                available_sampling.append('nearest')
+                if dim == 1:
+                    available_sampling.append('linear')
+                elif dim == 2:
+                    available_sampling.append('bilinear')
+                    available_sampling.append('bicubic')
+                else:
+                    available_sampling.append('trilinear')
+            assert mode in available_sampling, \
+                '`mode` {} '.format(mode) + \
+                'is not one of the available sample ' + \
+                'modes {}.'.format(available_sampling)
+
+        self.conv_block = nn.ModuleList()
+
+        while len(self.conv_block) < num_layers:
+            use_up = up and not self.conv_block
+            self.conv_block.append(_get_layer(
+                ConvUpLayer if use_up else ConvLayer, layer_kwargs, wrap=True, noise=noise))
+            layer_kwargs.update(in_channels=out_channels)
+
+        self.projection = None
+        if resnet:
+            projection_kwargs = {
+                **layer_kwargs,
+                'in_channels': in_channels,
+                'kernel_size': 1,
+                'modulate': False,
+                'demodulate': False
+            }
+            self.projection = _get_layer(
+                ConvUpLayer if up else ConvLayer, projection_kwargs, wrap=False)
+
+        self.res_scale = 1 / np.sqrt(2)
+
+    def __len__(self):
+        """
+        Get the number of conv layers in this block.
+        """
+        return len(self.conv_block)
+
+    def forward(self, input, latents=None, **kwargs):
+        """
+        Run some input through this block and return the output.
+        Arguments:
+            input (torch.Tensor)
+            latents (torch.Tensor)
+        Returns:
+            output (torch.Tensor)
+        """
+        if latents.dim() == 2:
+            latents.unsqueeze(1)
+        if latents.size(1) == 1:
+            latents = latents.repeat(1, len(self), 1)
+        assert latents.size(1) == len(self), \
+            'Number of latent inputs ' + \
+            '({}) does not match '.format(latents.size(1)) + \
+            'number of conv layers ' + \
+            '({}) in block.'.format(len(self))
+        x = input
+        for i, layer in enumerate(self.conv_block):
+            x = layer(input=x, latent=latents[:, i])
+        if self.projection is not None:
+            x += self.projection(input=input)
+            x *= self.res_scale
+        return x
+
+
+class DiscriminatorConvBlock(nn.Module):
+    """
+    A convblock for the discriminator model.
+    Arguments:
+        in_channels (int): Number of input channels.
+        out_channels (int): Number of output channels.
+        demodulate (bool): Normalize feature outputs from conv
+            layers. Default value is True.
+        resnet (bool): Use residual connections. Default value is
+            False.
+        down (bool): Downsample the data to twice its size. This is
+            performed in the last layer of the block. Default
+            value is False.
+        num_layers (int): Number of convolutional layers of this
+            block. Default value is 2.
+        filter (int, list): The filter to use if
+            `down=True` and `mode='FIR'`. If int, a low
+            pass filter of this size will be used. If list,
+            the filter is explicitly specified. If the filter
+            is of a single dimension it will be expanded to
+            the number of dimensions of the data. Default
+            value is a low pass filter of [1, 3, 3, 1].
+        activation (str, callable, nn.Module): The non-linear
+            activation function to use.
+            Default value is leaky relu with a slope of 0.2.
+        mode (str): The resample mode of downsampling layers.
+            Only used when `down=True`. If fused=True` only 'FIR'
+            and 'none' can be used. Else, 'max' or anything that can
+            be passed to torch.nn.functional.interpolate is
+            a valid mode. Default value is 'FIR'.
+        fused (bool): If `down=True`, fuse the downsample operation
+            and the last convolutional layer into a strided
+            convolutional layer.
+        kernel_size (int): Size of the convolutional kernel.
+            Default value is 3.
+        pad_mode (str): The padding mode for convolutional
+            layers. Has to be one of 'constant', 'reflect',
+            'replicate' or 'circular'. Default value is
+            'constant'.
+        pad_constant (float): The value to use for conv
+            padding if `conv_pad_mode='constant'`. Default
+            value is 0.
+        filter_pad_mode (str): If `mode='FIR'`, this is used with the filter.
+            Otherwise works the same as `pad_mode`.
+        filter_pad_constant (float): If `mode='FIR'`, this is used with the filter.
+            Otherwise works the same as `pad_constant`
+        pad_once (bool): If FIR filter is used, do all the padding for
+            both convolution and FIR in the FIR layer instead of once per layer.
+            Default value is True.
+        use_bias (bool): Add bias to layer outputs. Default value is True.
+        lr_mul (float): The learning rate multiplier for this
+            block. When loading weights of previously trained
+            networks, this value has to be the same as when
+            the network was trained for the outputs to not
+            change (as this is used to scale the weights).
+            Default value is 1.
+        weight_scale (bool): Use weight scaling for
+            equalized learning rate. Default value
+            is True.
+    """
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 resnet=False,
+                 down=False,
+                 num_layers=2,
+                 filter=[1, 3, 3, 1],
+                 activation='leaky:0.2',
+                 mode='FIR',
+                 fused=True,
+                 kernel_size=3,
+                 pad_mode='constant',
+                 pad_constant=0,
+                 filter_pad_mode='constant',
+                 filter_pad_constant=0,
+                 pad_once=True,
+                 use_bias=True,
+                 lr_mul=1,
+                 weight_scale=True,
+                 gain=1,
+                 dim=2,
+                 *args,
+                 **kwargs):
+        super(DiscriminatorConvBlock, self).__init__()
+        layer_kwargs = locals()
+        layer_kwargs.pop('self')
+        layer_kwargs.pop('__class__')
+        layer_kwargs.update(
+            out_channels=in_channels,
+            features=in_channels,
+            modulate=False,
+            demodulate=False
+        )
+
+        assert num_layers > 0
+        assert 1 <= dim <= 3, '`dim` can only be 1, 2 or 3.'
+        if down:
+            available_sampling = ['FIR']
+            if fused:
+                available_sampling.append('none')
+            else:
+                available_sampling.append('max')
+                available_sampling.append('area')
+                available_sampling.append('nearest')
+                if dim == 1:
+                    available_sampling.append('linear')
+                elif dim == 2:
+                    available_sampling.append('bilinear')
+                    available_sampling.append('bicubic')
+                else:
+                    available_sampling.append('trilinear')
+            assert mode in available_sampling, \
+                '`mode` {} '.format(mode) + \
+                'is not one of the available sample ' + \
+                'modes {}'.format(available_sampling)
+
+        self.conv_block = nn.ModuleList()
+
+        while len(self.conv_block) < num_layers:
+            if len(self.conv_block) == num_layers - 1:
+                layer_kwargs.update(
+                    out_channels=out_channels,
+                    features=out_channels
+                )
+            use_down = down and len(self.conv_block) == num_layers - 1
+            self.conv_block.append(_get_layer(
+                ConvDownLayer if use_down else ConvLayer, layer_kwargs, wrap=True, noise=False))
+
+        self.projection = None
+        if resnet:
+            projection_kwargs = {
+                **layer_kwargs,
+                'in_channels': in_channels,
+                'kernel_size': 1,
+                'modulate': False,
+                'demodulate': False
+            }
+            self.projection = _get_layer(
+                ConvDownLayer if down else ConvLayer, projection_kwargs, wrap=False)
+
+        self.res_scale = 1 / np.sqrt(2)
+
+    def __len__(self):
+        """
+        Get the number of conv layers in this block.
+        """
+        return len(self.conv_block)
+
+    def forward(self, input, **kwargs):
+        """
+        Run some input through this block and return the output.
+        Arguments:
+            input (torch.Tensor)
+        Returns:
+            output (torch.Tensor)
+        """
+        x = input
+        for layer in self.conv_block:
+            x = layer(input=x)
+        if self.projection is not None:
+            x += self.projection(input=input)
+            x *= self.res_scale
+        return x

stylegan2/project.py

@@ -0,0 +1,304 @@
+import warnings
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from . import models, utils
+from .external_models import lpips
+
+
+class Projector(nn.Module):
+    """
+    Projects data to latent space and noise tensors.
+    Arguments:
+        G (Generator)
+        dlatent_avg_samples (int): Number of dlatent samples
+            to collect to find the mean and std.
+            Default value is 10 000.
+        dlatent_avg_label (int, torch.Tensor, optional): The label to
+            use when gathering dlatent statistics.
+        dlatent_device (int, str, torch.device, optional): Device to use
+            for gathering statistics of dlatents. By default uses
+            the same device as parameters of `G` reside on.
+        dlatent_batch_size (int): The batch size to sample
+            dlatents with. Default value is 1024.
+        lpips_model (nn.Module): A model that returns feature the distance
+            between two inputs. Default value is the LPIPS VGG16 model.
+        lpips_size (int, optional): Resize any data fed to `lpips_model` by scaling
+            the data so that its smallest side is the same size as this
+            argument. Only has a default value of 256 if `lpips_model` is unspecified.
+        verbose (bool): Write progress of dlatent statistics gathering to stdout.
+            Default value is True.
+    """
+    def __init__(self,
+                 G,
+                 dlatent_avg_samples=10000,
+                 dlatent_avg_label=None,
+                 dlatent_device=None,
+                 dlatent_batch_size=1024,
+                 lpips_model=None,
+                 lpips_size=None,
+                 verbose=True):
+        super(Projector, self).__init__()
+        assert isinstance(G, models.Generator)
+        G.eval().requires_grad_(False)
+
+        self.G_synthesis = G.G_synthesis
+
+        G_mapping = G.G_mapping
+
+        dlatent_batch_size = min(dlatent_batch_size, dlatent_avg_samples)
+
+        if dlatent_device is None:
+            dlatent_device = next(G_mapping.parameters()).device()
+        else:
+            dlatent_device = torch.device(dlatent_device)
+
+        G_mapping.to(dlatent_device)
+
+        latents = torch.empty(
+            dlatent_avg_samples, G_mapping.latent_size).normal_()
+        dlatents = []
+
+        labels = None
+        if dlatent_avg_label is not None:
+            labels = torch.tensor(dlatent_avg_label).to(dlatent_device).long().view(-1).repeat(dlatent_batch_size)
+
+        if verbose:
+            progress = utils.ProgressWriter(np.ceil(dlatent_avg_samples / dlatent_batch_size))
+            progress.write('Gathering dlatents...', step=False)
+
+        for i in range(0, dlatent_avg_samples, dlatent_batch_size):
+            batch_latents = latents[i: i + dlatent_batch_size].to(dlatent_device)
+            batch_labels = None
+            if labels is not None:
+                batch_labels = labels[:len(batch_latents)]
+            with torch.no_grad():
+                dlatents.append(G_mapping(batch_latents, labels=batch_labels).cpu())
+            if verbose:
+                progress.step()
+
+        if verbose:
+            progress.write('Done!', step=False)
+            progress.close()
+
+        dlatents = torch.cat(dlatents, dim=0)
+
+        self.register_buffer(
+            '_dlatent_avg',
+            dlatents.mean(dim=0).view(1, 1, -1)
+        )
+        self.register_buffer(
+            '_dlatent_std',
+            torch.sqrt(
+                torch.sum((dlatents - self._dlatent_avg) ** 2) / dlatent_avg_samples + 1e-8
+            ).view(1, 1, 1)
+        )
+
+        if lpips_model is None:
+            warnings.warn(
+                'Using default LPIPS distance metric based on VGG 16. ' + \
+                'This metric will only work on image data where values are in ' + \
+                'the range [-1, 1], please specify an lpips module if you want ' + \
+                'to use other kinds of data formats.'
+            )
+            lpips_model = lpips.LPIPS_VGG16(pixel_min=-1, pixel_max=1)
+            lpips_size = 256
+        self.lpips_model = lpips_model.eval().requires_grad_(False)
+        self.lpips_size = lpips_size
+
+        self.to(dlatent_device)
+
+    def _scale_for_lpips(self, data):
+        if not self.lpips_size:
+            return data
+        scale_factor = self.lpips_size / min(data.size()[2:])
+        if scale_factor == 1:
+            return data
+        mode = 'nearest'
+        if scale_factor < 1:
+            mode = 'area'
+        return F.interpolate(data, scale_factor=scale_factor, mode=mode)
+
+    def _check_job(self):
+        assert self._job is not None, 'Call `start()` first to set up target.'
+        # device of dlatent param will not change with the rest of the models
+        # and buffers of this class as it was never registered as a buffer or
+        # parameter. Same goes for optimizer. Make sure it is on the correct device.
+        if self._job.dlatent_param.device != self._dlatent_avg.device:
+            self._job.dlatent_param = self._job.dlatent_param.to(self._dlatent_avg)
+            self._job.opt.load_state_dict(
+                utils.move_to_device(self._job.opt.state_dict(), self._dlatent_avg.device)[0])
+
+    def generate(self):
+        """
+        Generate an output with the current dlatent and noise values.
+        Returns:
+            output (torch.Tensor)
+        """
+        self._check_job()
+        with torch.no_grad():
+            return self.G_synthesis(self._job.dlatent_param)
+
+    def get_dlatent(self):
+        """
+        Get a copy of the current dlatent values.
+        Returns:
+            dlatents (torch.Tensor)
+        """
+        self._check_job()
+        return self._job.dlatent_param.data.clone()
+
+    def get_noise(self):
+        """
+        Get a copy of the current noise values.
+        Returns:
+            noise_tensors (list)
+        """
+        self._check_job()
+        return [noise.data.clone() for noise in self._job.noise_params]
+
+    def start(self,
+              target,
+              num_steps=1000,
+              initial_learning_rate=0.1,
+              initial_noise_factor=0.05,
+              lr_rampdown_length=0.25,
+              lr_rampup_length=0.05,
+              noise_ramp_length=0.75,
+              regularize_noise_weight=1e5,
+              verbose=True,
+              verbose_prefix=''):
+        """
+        Set up a target and its projection parameters.
+        Arguments:
+            target (torch.Tensor): The data target. This should
+                already be preprocessed (scaled to correct value range).
+            num_steps (int): Number of optimization steps. Default
+                value is 1000.
+            initial_learning_rate (float): Default value is 0.1.
+            initial_noise_factor (float): Default value is 0.05.
+            lr_rampdown_length (float): Default value is 0.25.
+            lr_rampup_length (float): Default value is 0.05.
+            noise_ramp_length (float): Default value is 0.75.
+            regularize_noise_weight (float): Default value is 1e5.
+            verbose (bool): Write progress to stdout every time
+                `step()` is called.
+            verbose_prefix (str, optional): This is written before
+                any other output to stdout.
+        """
+        if target.dim() == self.G_synthesis.dim + 1:
+            target = target.unsqueeze(0)
+        assert target.dim() == self.G_synthesis.dim + 2, \
+            'Number of dimensions of target data is incorrect.'
+
+        target = target.to(self._dlatent_avg)
+        target_scaled = self._scale_for_lpips(target)
+
+        dlatent_param = nn.Parameter(
+            self._dlatent_avg.clone().repeat(target.size(0), len(self.G_synthesis), 1))
+        noise_params = self.G_synthesis.static_noise(trainable=True)
+        params = [dlatent_param] + noise_params
+
+        opt = torch.optim.Adam(params)
+
+        noise_tensor = torch.empty_like(dlatent_param)
+
+        if verbose:
+            progress = utils.ProgressWriter(num_steps)
+            value_tracker = utils.ValueTracker()
+
+        self._job = utils.AttributeDict(**locals())
+        self._job.current_step = 0
+
+    def step(self, steps=1):
+        """
+        Take a projection step.
+        Arguments:
+            steps (int): Number of steps to take. If this
+                exceeds the remaining steps of the projection
+                that amount of steps is taken instead. Default
+                value is 1.
+        """
+        self._check_job()
+
+        remaining_steps = self._job.num_steps - self._job.current_step
+        if not remaining_steps > 0:
+            warnings.warn(
+                'Trying to take a projection step after the ' + \
+                'final projection iteration has been completed.'
+            )
+        if steps < 0:
+            steps = remaining_steps
+        steps = min(remaining_steps, steps)
+
+        if not steps > 0:
+            return
+
+        for _ in range(steps):
+
+            if self._job.current_step >= self._job.num_steps:
+                break
+
+            # Hyperparameters.
+            t = self._job.current_step / self._job.num_steps
+            noise_strength = self._dlatent_std * self._job.initial_noise_factor \
+                             * max(0.0, 1.0 - t / self._job.noise_ramp_length) ** 2
+            lr_ramp = min(1.0, (1.0 - t) / self._job.lr_rampdown_length)
+            lr_ramp = 0.5 - 0.5 * np.cos(lr_ramp * np.pi)
+            lr_ramp = lr_ramp * min(1.0, t / self._job.lr_rampup_length)
+            learning_rate = self._job.initial_learning_rate * lr_ramp
+
+            for param_group in self._job.opt.param_groups:
+                param_group['lr'] = learning_rate
+
+            dlatents = self._job.dlatent_param + noise_strength * self._job.noise_tensor.normal_()
+
+            output = self.G_synthesis(dlatents)
+            assert output.size() == self._job.target.size(), \
+                'target size {} does not fit output size {} of generator'.format(
+                    target.size(), output.size())
+
+            output_scaled = self._scale_for_lpips(output)
+
+            # Main loss: LPIPS distance of output and target
+            lpips_distance = torch.mean(self.lpips_model(output_scaled, self._job.target_scaled))
+
+            # Calculate noise regularization loss
+            reg_loss = 0
+            for p in self._job.noise_params:
+                size = min(p.size()[2:])
+                dim = p.dim() - 2
+                while True:
+                    reg_loss += torch.mean(
+                        (p * p.roll(shifts=[1] * dim, dims=list(range(2, 2 + dim)))) ** 2)
+                    if size <= 8:
+                        break
+                    p = F.interpolate(p, scale_factor=0.5, mode='area')
+                    size = size // 2
+
+            # Combine loss, backward and update params
+            loss = lpips_distance + self._job.regularize_noise_weight * reg_loss
+            self._job.opt.zero_grad()
+            loss.backward()
+            self._job.opt.step()
+
+            # Normalize noise values
+            for p in self._job.noise_params:
+                with torch.no_grad():
+                    p_mean = p.mean(dim=list(range(1, p.dim())), keepdim=True)
+                    p_rstd = torch.rsqrt(
+                        torch.mean((p - p_mean) ** 2, dim=list(range(1, p.dim())), keepdim=True) + 1e-8)
+                    p.data = (p.data - p_mean) * p_rstd
+
+            self._job.current_step += 1
+
+            if self._job.verbose:
+                self._job.value_tracker.add('loss', float(loss))
+                self._job.value_tracker.add('lpips_distance', float(lpips_distance))
+                self._job.value_tracker.add('noise_reg', float(reg_loss))
+                self._job.value_tracker.add('lr', learning_rate, beta=0)
+                self._job.progress.write(self._job.verbose_prefix, str(self._job.value_tracker))
+                if self._job.current_step >= self._job.num_steps:
+                    self._job.progress.close()

stylegan2/train.py

@@ -0,0 +1,1013 @@
+import warnings
+import functools
+import os
+import time
+import sys
+import json
+import numpy as np
+import torch
+import torch.utils.tensorboard
+from torch import nn
+import torchvision
+try:
+    import apex
+    from apex import amp
+except ImportError:
+    pass
+
+from . import models, utils, loss_fns
+
+
+class Trainer:
+    """
+    Class that handles training and logging for stylegan2.
+    For distributed training, the arguments `rank`, `world_size`,
+    `master_addr`, `master_port` can all be given as environmnet variables
+    (only difference is that the keys should be capital cased).
+    Environment variables if available will override any python
+    value for the same argument.
+    Arguments:
+        G (Generator): The generator model.
+        D (Discriminator): The discriminator model.
+        latent_size (int): The size of the latent inputs.
+        dataset (indexable object): The dataset. Has to implement
+            '__getitem__' and '__len__'. If `label_size` > 0, this
+            dataset object has to return both a data entry and its
+            label when calling '__getitem__'.
+        device (str, int, list, torch.device): The device to run training on.
+            Can be a list of integers for parallel training in the same
+            process. Parallel training can also be achieved by spawning
+            seperate processes and using the `rank` argument for each
+            process. In that case, only one device should be specified
+            per process.
+        Gs (Generator, optional): A generator copy with the current
+            moving average of the training generator. If not specified,
+            a copy of the generator is made for the moving average of
+            weights.
+        Gs_beta (float): The beta value for the moving average weights.
+            Default value is 1 / (2 ^(32 / 10000)).
+        Gs_device (str, int, torch.device, optional): The device to store
+            the moving average weights on. If using a different device
+            than what is specified for the `device` argument, updating
+            the moving average weights will take longer as the data
+            will have to be transfered over different devices. If
+            this argument is not specified, the same device is used
+            as specified in the `device` argument.
+        batch_size (int): The total batch size to average gradients
+            over. This should be the combined batch size of all used
+            devices (it is later divided by world size for distributed
+            training).
+            Example: We want to average gradients over 32 data
+                entries. To do this we just set `batch_size=32`.
+                Even if we train on 8 GPUs we still use the same
+                batch size (each GPU will take 4 data entries per
+                batch).
+            Default value is 32.
+        device_batch_size (int): The number of data entries that can
+            fit on the specified device at a time.
+            Example: We want to average gradients over 32 data
+                entries. To do this we just set `batch_size=32`.
+                However, our device can only handle a batch of
+                4 at a time before running out of memory. We
+                therefor set `device_batch_size=4`. With a
+                single device (no distributed training), each
+                batch is split into 32 / 4 parts and gradients
+                are averaged over all these parts.
+            Default value is 4.
+        label_size (int, optional): Number of possible class labels.
+            This is required for conditioning the GAN with labels.
+            If not specified it is assumed that no labels are used.
+        data_workers (int): The number of spawned processes that
+            handle data loading. Default value is 4.
+        G_loss (str, callable): The loss function to use
+            for the generator. If string, it can be one of the
+            following: 'logistic', 'logistic_ns' or 'wgan'.
+            If not a string, the callable has to follow
+            the format of functions found in `stylegan2.loss`.
+            Default value is 'logistic_ns' (non-saturating logistic).
+        D_loss (str, callable): The loss function to use
+            for the discriminator. If string, it can be one of the
+            following: 'logistic' or 'wgan'.
+            If not a string, same restriction follows as for `G_loss`.
+            Default value is 'logistic'.
+        G_reg (str, callable, None): The regularizer function to use
+            for the generator. If string, it can only be 'pathreg'
+            (pathlength regularization). A weight for the regularizer
+            can be passed after the string name like the following:
+                G_reg='pathreg:5'
+            This will assign a weight of 5 to the regularization loss.
+            If set to None, no geenerator regularization is performed.
+            Default value is 'pathreg:2'.
+        G_reg_interval (int): The interval at which to regularize the
+            generator. If set to 0, the regularization and loss gradients
+            are combined in a single optimization step every iteration.
+            If set to 1, the gradients for the regularization and loss
+            are used separately for two optimization steps. Any value
+            higher than 1 indicates that regularization should only
+            be performed at this interval (lazy regularization).
+            Default value is 4.
+        G_opt_class (str, class): The optimizer class for the generator.
+            Default value is 'Adam'.
+        G_opt_kwargs (dict): Keyword arguments for the generator optimizer
+            constructor. Default value is {'lr': 2e-3, 'betas': (0, 0.99)}.
+        G_reg_batch_size (int): Same as `batch_size` but only for
+            the regularization loss of the generator. Default value
+            is 16.
+        G_reg_device_batch_size (int): Same as `device_batch_size`
+            but only for the regularization loss of the generator.
+            Default value is 2.
+        D_reg (str, callable, None): The regularizer function to use
+            for the discriminator. If string, the following values
+            can be used: 'r1', 'r2', 'gp'. See doc for `G_reg` for
+            rest of info on regularizer format.
+            Default value is 'r1:10'.
+        D_reg_interval (int): Same as `D_reg_interval` but for the
+            discriminator. Default value is 16.
+        D_opt_class (str, class): The optimizer class for the discriminator.
+            Default value is 'Adam'.
+        D_opt_kwargs (dict): Keyword arguments for the discriminator optimizer
+            constructor. Default value is {'lr': 2e-3, 'betas': (0, 0.99)}.
+        style_mix_prob (float): The probability of passing 2 latents instead of 1
+            to the generator during training. Default value is 0.9.
+        G_iter (int): Number of generator iterations for every full training
+            iteration. Default value is 1.
+        D_iter (int): Number of discriminator iterations for every full training
+            iteration. Default value is 1.
+        pl_avg (float, torch.Tensor): The average pathlength starting value for
+            pathlength regularization of the generator. Default value is 0.
+        tensorboard_log_dir (str, optional): A path to a directory to log training values
+            in for tensorboard. Only used without distributed training or when
+            distributed training is enabled and the rank of this trainer is 0.
+        checkpoint_dir (str, optional): A path to a directory to save training
+            checkpoints to. If not specified, not checkpoints are automatically
+            saved during training.
+        checkpoint_interval (int): The interval at which to save training checkpoints.
+            Default value is 10000.
+        seen (int): The number of previously trained iterations. Used for logging.
+            Default value is 0.
+        half (bool): Use mixed precision training. Default value is False.
+        rank (int, optional): If set, use distributed training. Expects that
+            this object has been constructed with the same arguments except
+            for `rank` in different processes.
+        world_size (int, optional): If using distributed training, this specifies
+            the number of nodes in the training.
+        master_addr (str): The master address for distributed training.
+            Default value is '127.0.0.1'.
+        master_port (str): The master port for distributed training.
+            Default value is '23456'.
+    """
+
+    def __init__(self,
+                 G,
+                 D,
+                 latent_size,
+                 dataset,
+                 device,
+                 Gs=None,
+                 Gs_beta=0.5 ** (32 / 10000),
+                 Gs_device=None,
+                 batch_size=32,
+                 device_batch_size=4,
+                 label_size=0,
+                 data_workers=4,
+                 G_loss='logistic_ns',
+                 D_loss='logistic',
+                 G_reg='pathreg:2',
+                 G_reg_interval=4,
+                 G_opt_class='Adam',
+                 G_opt_kwargs={'lr': 2e-3, 'betas': (0, 0.99)},
+                 G_reg_batch_size=None,
+                 G_reg_device_batch_size=None,
+                 D_reg='r1:10',
+                 D_reg_interval=16,
+                 D_opt_class='Adam',
+                 D_opt_kwargs={'lr': 2e-3, 'betas': (0, 0.99)},
+                 style_mix_prob=0.9,
+                 G_iter=1,
+                 D_iter=1,
+                 pl_avg=0.,
+                 tensorboard_log_dir=None,
+                 checkpoint_dir=None,
+                 checkpoint_interval=10000,
+                 seen=0,
+                 half=False,
+                 rank=None,
+                 world_size=None,
+                 master_addr='127.0.0.1',
+                 master_port='23456'):
+        assert not isinstance(G, nn.parallel.DistributedDataParallel) and \
+               not isinstance(D, nn.parallel.DistributedDataParallel), \
+               'Encountered a model wrapped in `DistributedDataParallel`. ' + \
+               'Distributed parallelism is handled by this class and can ' + \
+               'not be initialized before.'
+        # We store the training settings in a dict that can be saved as a json file.
+        kwargs = locals()
+        # First we remove the arguments that can not be turned into json.
+        kwargs.pop('self')
+        kwargs.pop('G')
+        kwargs.pop('D')
+        kwargs.pop('Gs')
+        kwargs.pop('dataset')
+        # Some arguments may have to be turned into strings to be compatible with json.
+        kwargs.update(pl_avg=float(pl_avg))
+        if isinstance(device, torch.device):
+            kwargs.update(device=str(device))
+        if isinstance(Gs_device, torch.device):
+            kwargs.update(device=str(Gs_device))
+        self.kwargs = kwargs
+
+        if device or device == 0:
+            if isinstance(device, (tuple, list)):
+                self.device = torch.device(device[0])
+            else:
+                self.device = torch.device(device)
+        else:
+            self.device = torch.device('cpu')
+        if self.device.index is not None:
+            torch.cuda.set_device(self.device.index)
+        else:
+            assert not half, 'Mixed precision training only available ' + \
+                'for CUDA devices.'
+        # Set up the models
+        self.G = G.train().to(self.device)
+        self.D = D.train().to(self.device)
+        if isinstance(device, (tuple, list)) and len(device) > 1:
+            assert all(isinstance(dev, int) for dev in device), \
+                'Multiple devices have to be specified as a list ' + \
+                'or tuple of integers corresponding to device indices.'
+            # TODO: Look into bug with torch.autograd.grad and nn.DataParallel
+            # In the meanwhile just prohibit its use together.
+            assert G_reg is None and D_reg is None, 'Regularization ' + \
+                'currently not supported for multi-gpu training in single process. ' + \
+                'Please use distributed training with one device per process instead.'
+
+            device_batch_size *= len(device)
+            def to_data_parallel(model):
+                if not isinstance(model, nn.DataParallel):
+                    return nn.DataParallel(model, device_ids=device)
+                return model
+            self.G = to_data_parallel(self.G)
+            self.D = to_data_parallel(self.D)
+
+        # Default generator reg batch size is the global batch size
+        # unless it has been specified otherwise.
+        G_reg_batch_size = G_reg_batch_size or batch_size
+        G_reg_device_batch_size = G_reg_device_batch_size or device_batch_size
+
+        # Set up distributed training
+        rank = os.environ.get('RANK', rank)
+        if rank is not None:
+            rank = int(rank)
+            addr = os.environ.get('MASTER_ADDR', master_addr)
+            port = os.environ.get('MASTER_PORT', master_port)
+            world_size = os.environ.get('WORLD_SIZE', world_size)
+            assert world_size is not None, 'Distributed training ' + \
+                'requires specifying world size.'
+            world_size = int(world_size)
+            assert self.device.index is not None, \
+                'Distributed training is only supported for CUDA.'
+            assert batch_size % world_size == 0, 'Batch size has to be ' + \
+                'evenly divisible by world size.'
+            assert G_reg_batch_size % world_size == 0, 'G reg batch size has to be ' + \
+                'evenly divisible by world size.'
+            batch_size = batch_size // world_size
+            G_reg_batch_size = G_reg_batch_size // world_size
+            init_method = 'tcp://{}:{}'.format(addr, port)
+            torch.distributed.init_process_group(
+                backend='nccl', init_method=init_method, rank=rank, world_size=world_size)
+        else:
+            world_size = 1
+        self.rank = rank
+        self.world_size = world_size
+
+        # Set up variable to keep track of moving average of path lengths
+        self.pl_avg = torch.tensor(
+            pl_avg, dtype=torch.float16 if half else torch.float32, device=self.device)
+
+        # Broadcast parameters from rank 0 if running distributed
+        self._sync_distributed(G=self.G, D=self.D, broadcast_weights=True)
+
+        # Set up moving average of generator
+        # Only for non-distributed training or
+        # if rank is 0
+        if not self.rank:
+            # Values for `rank`: None -> not distributed, 0 -> distributed and 'main' node
+            self.Gs = Gs
+            if not isinstance(Gs, utils.MovingAverageModule):
+                self.Gs = utils.MovingAverageModule(
+                    from_module=self.G,
+                    to_module=Gs,
+                    param_beta=Gs_beta,
+                    device=self.device if Gs_device is None else Gs_device
+                )
+        else:
+            self.Gs = None
+
+        # Set up loss and regularization functions
+        self.G_loss = get_loss_fn('G', G_loss)
+        self.D_loss = get_loss_fn('D', D_loss)
+        self.G_reg = get_reg_fn('G', G_reg, pl_avg=self.pl_avg)
+        self.D_reg = get_reg_fn('D', D_reg)
+        self.G_reg_interval = G_reg_interval
+        self.D_reg_interval = D_reg_interval
+        self.G_iter = G_iter
+        self.D_iter = D_iter
+
+        # Set up optimizers (adjust hyperparameters if lazy regularization is active)
+        self.G_opt = build_opt(self.G, G_opt_class, G_opt_kwargs, self.G_reg, self.G_reg_interval)
+        self.D_opt = build_opt(self.D, D_opt_class, D_opt_kwargs, self.D_reg, self.D_reg_interval)
+
+        # Set up mixed precision training
+        if half:
+            assert 'apex' in sys.modules, 'Can not run mixed precision ' + \
+                'training (`half=True`) without the apex module.'
+            (self.G, self.D), (self.G_opt, self.D_opt) = amp.initialize(
+                [self.G, self.D], [self.G_opt, self.D_opt], opt_level='O1')
+        self.half = half
+
+        # Data
+        sampler = None
+        if self.rank is not None:
+            sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=True)
+        self.dataloader = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=device_batch_size,
+            num_workers=data_workers,
+            shuffle=sampler is None,
+            pin_memory=self.device.index is not None,
+            drop_last=True,
+            sampler=sampler
+        )
+        self.dataloader_iter = None
+        self.prior_generator = utils.PriorGenerator(
+            latent_size=latent_size,
+            label_size=label_size,
+            batch_size=device_batch_size,
+            device=self.device
+        )
+        assert batch_size % device_batch_size == 0, \
+            'Batch size has to be evenly divisible by the product of ' + \
+            'device batch size and world size.'
+        self.subdivisions = batch_size // device_batch_size
+        assert G_reg_batch_size % G_reg_device_batch_size == 0, \
+            'G reg batch size has to be evenly divisible by the product of ' + \
+            'G reg device batch size and world size.'
+        self.G_reg_subdivisions = G_reg_batch_size // G_reg_device_batch_size
+        self.G_reg_device_batch_size = G_reg_device_batch_size
+
+        self.tb_writer = None
+        if tensorboard_log_dir and not self.rank:
+            self.tb_writer = torch.utils.tensorboard.SummaryWriter(tensorboard_log_dir)
+
+        self.label_size = label_size
+        self.style_mix_prob = style_mix_prob
+        self.checkpoint_dir = checkpoint_dir
+        self.checkpoint_interval = checkpoint_interval
+        self.seen = seen
+        self.metrics = {}
+        self.callbacks = []
+
+    def _get_batch(self):
+        """
+        Fetch a batch and its labels. If no labels are
+        available the returned labels will be `None`.
+        Returns:
+            data
+            labels
+        """
+        if self.dataloader_iter is None:
+            self.dataloader_iter = iter(self.dataloader)
+        try:
+            batch = next(self.dataloader_iter)
+        except StopIteration:
+            self.dataloader_iter = None
+            return self._get_batch()
+        if isinstance(batch, (tuple, list)):
+            if len(batch) > 1:
+                data, label = batch[:2]
+            else:
+                data, label = batch[0], None
+        else:
+            data, label = batch, None
+        if not self.label_size:
+            label = None
+        if torch.is_tensor(data):
+            data = data.to(self.device)
+        if torch.is_tensor(label):
+            label = label.to(self.device)
+        return data, label
+
+    def _sync_distributed(self, G=None, D=None, broadcast_weights=False):
+        """
+        Sync the gradients (and alternatively the weights) of
+        the specified networks over the distributed training
+        nodes. Varying buffers are broadcasted from rank 0.
+        If no distributed training is not enabled, no action
+        is taken and this is a no-op function.
+        Arguments:
+            G (Generator, optional)
+            D (Discriminator, optional)
+            broadcast_weights (bool): Broadcast the weights from
+                node of rank 0 to all other ranks. Default
+                value is False.
+        """
+        if self.rank is None:
+            return
+        for net in [G, D]:
+            if net is None:
+                continue
+            for p in net.parameters():
+                if p.grad is not None:
+                    torch.distributed.all_reduce(p.grad, async_op=True)
+                if broadcast_weights:
+                    torch.distributed.broadcast(p.data, src=0, async_op=True)
+        if G is not None:
+            if G.dlatent_avg is not None:
+                torch.distributed.broadcast(G.dlatent_avg, src=0, async_op=True)
+            if self.pl_avg is not None:
+                torch.distributed.broadcast(self.pl_avg, src=0, async_op=True)
+        if G is not None or D is not None:
+            torch.distributed.barrier(async_op=False)
+
+    def _backward(self, loss, opt, mul=1, subdivisions=None):
+        """
+        Reduce loss by world size and subdivisions before
+        calling backward for the loss. Loss scaling is
+        performed when mixed precision training is
+        enabled.
+        Arguments:
+            loss (torch.Tensor)
+            opt (torch.optim.Optimizer)
+            mul (float): Loss weight. Default value is 1.
+            subdivisions (int, optional): The number of
+                subdivisions to divide by. If this is
+                not specified, the subdvisions from
+                the specified batch and device size
+                at construction is used.
+        Returns:
+            loss (torch.Tensor): The loss scaled by mul
+                and subdivisions but not by world size.
+        """
+        if loss is None:
+            return 0
+        mul /= subdivisions or self.subdivisions
+        mul /= self.world_size or 1
+        if mul != 1:
+            loss *= mul
+        if self.half:
+            with amp.scale_loss(loss, opt) as scaled_loss:
+                scaled_loss.backward()
+        else:
+            loss.backward()
+        #get the scalar only
+        return loss.item() * (self.world_size or 1)
+
+    def train(self, iterations, callbacks=None, verbose=True):
+        """
+        Train the models for a specific number of iterations.
+        Arguments:
+            iterations (int): Number of iterations to train for.
+            callbacks (callable, list, optional): One
+                or more callbacks to call at the end of each
+                iteration. The function is given the total
+                number of batches that have been processed since
+                this trainer object was initialized (not reset
+                when loading a saved checkpoint).
+                Default value is None (unused).
+            verbose (bool): Write progress to stdout.
+                Default value is True.
+        """
+        evaluated_metrics = {}
+        if self.rank:
+            verbose=False
+        if verbose:
+            progress = utils.ProgressWriter(iterations)
+            value_tracker = utils.ValueTracker()
+        for _ in range(iterations):
+            # Figure out if G and/or D be
+            # regularized this iteration
+            G_reg = self.G_reg is not None
+            if self.G_reg_interval and G_reg:
+                G_reg = self.seen % self.G_reg_interval == 0
+            D_reg = self.D_reg is not None
+            if self.D_reg_interval and D_reg:
+                D_reg = self.seen % self.D_reg_interval == 0
+
+            # -----| Train G |----- #
+
+            # Disable gradients for D while training G
+            self.D.requires_grad_(False)
+
+            for _ in range(self.G_iter):
+                self.G_opt.zero_grad()
+
+                G_loss = 0
+                for i in range(self.subdivisions):
+                    latents, latent_labels = self.prior_generator(
+                        multi_latent_prob=self.style_mix_prob)
+                    loss, _ = self.G_loss(
+                        G=self.G,
+                        D=self.D,
+                        latents=latents,
+                        latent_labels=latent_labels
+                    )
+                    G_loss += self._backward(loss, self.G_opt)
+
+                if G_reg:
+                    if self.G_reg_interval:
+                        # For lazy regularization, even if the interval
+                        # is set to 1, the optimization step is taken
+                        # before the gradients of the regularization is gathered.
+                        self._sync_distributed(G=self.G)
+                        self.G_opt.step()
+                        self.G_opt.zero_grad()
+                    G_reg_loss = 0
+                    # Pathreg is expensive to compute which
+                    # is why G regularization has its own settings
+                    # for subdivisions and batch size.
+                    for i in range(self.G_reg_subdivisions):
+                        latents, latent_labels = self.prior_generator(
+                            batch_size=self.G_reg_device_batch_size,
+                            multi_latent_prob=self.style_mix_prob
+                        )
+                        _, reg_loss = self.G_reg(
+                            G=self.G,
+                            latents=latents,
+                            latent_labels=latent_labels
+                        )
+                        G_reg_loss += self._backward(
+                            reg_loss,
+                            self.G_opt, mul=self.G_reg_interval or 1,
+                            subdivisions=self.G_reg_subdivisions
+                        )
+                self._sync_distributed(G=self.G)
+                self.G_opt.step()
+                # Update moving average of weights after
+                # each G training subiteration
+                if self.Gs is not None:
+                    self.Gs.update()
+
+            # Re-enable gradients for D
+            self.D.requires_grad_(True)
+
+            # -----| Train D |----- #
+
+            # Disable gradients for G while training D
+            self.G.requires_grad_(False)
+
+            for _ in range(self.D_iter):
+                self.D_opt.zero_grad()
+
+                D_loss = 0
+                for i in range(self.subdivisions):
+                    latents, latent_labels = self.prior_generator(
+                        multi_latent_prob=self.style_mix_prob)
+                    reals, real_labels = self._get_batch()
+                    loss, _ = self.D_loss(
+                        G=self.G,
+                        D=self.D,
+                        latents=latents,
+                        latent_labels=latent_labels,
+                        reals=reals,
+                        real_labels=real_labels
+                    )
+                    D_loss += self._backward(loss, self.D_opt)
+
+                if D_reg:
+                    if self.D_reg_interval:
+                        # For lazy regularization, even if the interval
+                        # is set to 1, the optimization step is taken
+                        # before the gradients of the regularization is gathered.
+                        self._sync_distributed(D=self.D)
+                        self.D_opt.step()
+                        self.D_opt.zero_grad()
+                    D_reg_loss = 0
+                    for i in range(self.subdivisions):
+                        latents, latent_labels = self.prior_generator(
+                            multi_latent_prob=self.style_mix_prob)
+                        reals, real_labels = self._get_batch()
+                        _, reg_loss = self.D_reg(
+                            G=self.G,
+                            D=self.D,
+                            latents=latents,
+                            latent_labels=latent_labels,
+                            reals=reals,
+                            real_labels=real_labels
+                        )
+                        D_reg_loss += self._backward(
+                            reg_loss, self.D_opt, mul=self.D_reg_interval or 1)
+                self._sync_distributed(D=self.D)
+                self.D_opt.step()
+
+            # Re-enable grads for G
+            self.G.requires_grad_(True)
+
+            if self.tb_writer is not None or verbose:
+                # In case verbose is true and tensorboard logging enabled
+                # we calculate grad norm here to only do it once as well
+                # as making sure we do it before any metrics that may
+                # possibly zero the grads.
+                G_grad_norm = utils.get_grad_norm_from_optimizer(self.G_opt)
+                D_grad_norm = utils.get_grad_norm_from_optimizer(self.D_opt)
+
+            for name, metric in self.metrics.items():
+                if not metric['interval'] or self.seen % metric['interval'] == 0:
+                    evaluated_metrics[name] = metric['eval_fn']()
+
+            # Printing and logging
+
+            # Tensorboard logging
+            if self.tb_writer is not None:
+                self.tb_writer.add_scalar('Loss/G_loss', G_loss, self.seen)
+                if G_reg:
+                    self.tb_writer.add_scalar('Loss/G_reg', G_reg_loss, self.seen)
+                    self.tb_writer.add_scalar('Grad_norm/G_reg', G_grad_norm, self.seen)
+                    self.tb_writer.add_scalar('Params/pl_avg', self.pl_avg, self.seen)
+                else:
+                    self.tb_writer.add_scalar('Grad_norm/G_loss', G_grad_norm, self.seen)
+                self.tb_writer.add_scalar('Loss/D_loss', D_loss, self.seen)
+                if D_reg:
+                    self.tb_writer.add_scalar('Loss/D_reg', D_reg_loss, self.seen)
+                    self.tb_writer.add_scalar('Grad_norm/D_reg', D_grad_norm, self.seen)
+                else:
+                    self.tb_writer.add_scalar('Grad_norm/D_loss', D_grad_norm, self.seen)
+                for name, value in evaluated_metrics.items():
+                    self.tb_writer.add_scalar('Metrics/{}'.format(name), value, self.seen)
+
+            # Printing
+            if verbose:
+                value_tracker.add('seen', self.seen + 1, beta=0)
+                value_tracker.add('G_lr', self.G_opt.param_groups[0]['lr'], beta=0)
+                value_tracker.add('G_loss', G_loss)
+                if G_reg:
+                    value_tracker.add('G_reg', G_reg_loss)
+                    value_tracker.add('G_reg_grad_norm', G_grad_norm)
+                    value_tracker.add('pl_avg', self.pl_avg, beta=0)
+                else:
+                    value_tracker.add('G_loss_grad_norm', G_grad_norm)
+                value_tracker.add('D_lr', self.D_opt.param_groups[0]['lr'], beta=0)
+                value_tracker.add('D_loss', D_loss)
+                if D_reg:
+                    value_tracker.add('D_reg', D_reg_loss)
+                    value_tracker.add('D_reg_grad_norm', D_grad_norm)
+                else:
+                    value_tracker.add('D_loss_grad_norm', D_grad_norm)
+                for name, value in evaluated_metrics.items():
+                    value_tracker.add(name, value, beta=0)
+                progress.write(str(value_tracker))
+
+            # Callback
+            for callback in utils.to_list(callbacks) + self.callbacks:
+                callback(self.seen)
+
+            self.seen += 1
+            
+            # clear cache
+            torch.cuda.empty_cache()
+            # Handle checkpointing
+            if not self.rank and self.checkpoint_dir and self.checkpoint_interval:
+                if self.seen % self.checkpoint_interval == 0:
+                    checkpoint_path = os.path.join(
+                        self.checkpoint_dir,
+                        '{}_{}'.format(self.seen, time.strftime('%Y-%m-%d_%H-%M-%S'))
+                    )
+                    self.save_checkpoint(checkpoint_path)
+
+        if verbose:
+            progress.close()
+
+    def register_metric(self, name, eval_fn, interval):
+        """
+        Add a metric. This will be evaluated every `interval`
+        training iteration. Used by tensorboard and progress
+        updates written to stdout while training.
+        Arguments:
+            name (str): A name for the metric. If a metric with
+                this name already exists it will be overwritten.
+            eval_fn (callable): A function that evaluates the metric
+                and returns a python number.
+            interval (int): The interval to evaluate at.
+        """
+        self.metrics[name] = {'eval_fn': eval_fn, 'interval': interval}
+
+    def remove_metric(self, name):
+        """
+        Remove a metric that was previously registered.
+        Arguments:
+            name (str): Name of the metric.
+        """
+        if name in self.metrics:
+            del self.metrics[name]
+        else:
+            warnings.warn(
+                'Attempting to remove metric {} '.format(name) + \
+                'which does not exist.'
+            )
+
+    def generate_images(self,
+                        num_images,
+                        seed=None,
+                        truncation_psi=None,
+                        truncation_cutoff=None,
+                        label=None,
+                        pixel_min=-1,
+                        pixel_max=1):
+        """
+        Generate some images with the generator and transform them into PIL
+        images and return them as a list.
+        Arguments:
+            num_images (int): Number of images to generate.
+            seed (int, optional): The seed for the random generation
+                of input latent values.
+            truncation_psi (float): See stylegan2.model.Generator.set_truncation()
+                Default value is None.
+            truncation_cutoff (int): See stylegan2.model.Generator.set_truncation()
+            label (int, list, optional): Label to condition all generated images with
+                or multiple labels, one for each generated image.
+            pixel_min (float): The min value in the pixel range of the generator.
+                Default value is -1.
+            pixel_min (float): The max value in the pixel range of the generator.
+                Default value is 1.
+        Returns:
+            images (list): List of PIL images.
+        """
+        if seed is None:
+            seed = int(10000 * time.time())
+        latents, latent_labels = self.prior_generator(num_images, seed=seed)
+        if label:
+            assert latent_labels is not None, 'Can not specify label when no labels ' + \
+                'are used by this model.'
+            label = utils.to_list(label)
+            assert all(isinstance(l, int) for l in label), '`label` can only consist of ' + \
+                'one or more python integers.'
+            assert len(label) == 1 or len(label) == num_images, '`label` can either ' + \
+                'specify one label to use for all images or a list of labels of the ' + \
+                'same length as number of images. Received {} labels '.format(len(label)) + \
+                'but {} images are to be generated.'.format(num_images)
+            if len(label) == 1:
+                latent_labels.fill_(label[0])
+            else:
+                latent_labels = torch.tensor(label).to(latent_labels)
+        self.Gs.set_truncation(
+            truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff)
+        with torch.no_grad():
+            generated = self.Gs(latents=latents, labels=latent_labels)
+        assert generated.dim() - 2 == 2, 'Can only generate images when using a ' + \
+            'network built for 2-dimensional data.'
+        assert generated.dim() == 4, 'Only generators that produce 2d data ' + \
+            'can be used to generate images.'
+        return utils.tensor_to_PIL(generated, pixel_min=pixel_min, pixel_max=pixel_max)
+
+    def log_images_tensorboard(self, images, name, resize=256):
+        """
+        Log a list of images to tensorboard by first turning
+        them into a grid. Can not be performed if rank > 0
+        or tensorboard_log_dir was not given at construction.
+        Arguments:
+            images (list): List of PIL images.
+            name (str): The name to log images for.
+            resize (int, tuple): The height and width to use for
+                each image in the grid. Default value is 256.
+        """
+        assert self.tb_writer is not None, \
+            'No tensorboard log dir was specified ' + \
+            'when constructing this object.'
+        image = utils.stack_images_PIL(images, individual_img_size=resize)
+        image = torchvision.transforms.ToTensor()(image)
+        self.tb_writer.add_image(name, image, self.seen)
+
+    def add_tensorboard_image_logging(self,
+                                      name,
+                                      interval,
+                                      num_images,
+                                      resize=256,
+                                      seed=None,
+                                      truncation_psi=None,
+                                      truncation_cutoff=None,
+                                      label=None,
+                                      pixel_min=-1,
+                                      pixel_max=1):
+        """
+        Set up tensorboard logging of generated images to be performed
+        at a certain training interval. If distributed training is set up
+        and this object does not have the rank 0, no logging will be performed
+        by this object.
+        All arguments except the ones mentioned below have their description
+        in the docstring of `generate_images()` and `log_images_tensorboard()`.
+        Arguments:
+            interval (int): The interval at which to log generated images.
+        """
+        if self.rank:
+            return
+        def callback(seen):
+            if seen % interval == 0:
+                images = self.generate_images(
+                    num_images=num_images,
+                    seed=seed,
+                    truncation_psi=truncation_psi,
+                    truncation_cutoff=truncation_cutoff,
+                    label=label,
+                    pixel_min=pixel_min,
+                    pixel_max=pixel_max
+                )
+                self.log_images_tensorboard(
+                    images=images,
+                    name=name,
+                    resize=resize
+                )
+        self.callbacks.append(callback)
+
+    def save_checkpoint(self, dir_path):
+        """
+        Save the current state of this trainer as a checkpoint.
+        NOTE: The dataset can not be serialized and saved so this
+            has to be reconstructed and given when loading this checkpoint.
+        Arguments:
+            dir_path (str): The checkpoint path.
+        """
+        if not os.path.exists(dir_path):
+            os.makedirs(dir_path)
+        else:
+            assert os.path.isdir(dir_path), '`dir_path` points to a file.'
+        kwargs = self.kwargs.copy()
+        # Update arguments that may have changed since construction
+        kwargs.update(
+            seen=self.seen,
+            pl_avg=float(self.pl_avg)
+        )
+        with open(os.path.join(dir_path, 'kwargs.json'), 'w') as fp:
+            json.dump(kwargs, fp)
+        torch.save(self.G_opt.state_dict(), os.path.join(dir_path, 'G_opt.pth'))
+        torch.save(self.D_opt.state_dict(), os.path.join(dir_path, 'D_opt.pth'))
+        models.save(self.G, os.path.join(dir_path, 'G.pth'))
+        models.save(self.D, os.path.join(dir_path, 'D.pth'))
+        if self.Gs is not None:
+            models.save(self.Gs, os.path.join(dir_path, 'Gs.pth'))
+
+    @classmethod
+    def load_checkpoint(cls, checkpoint_path, dataset, **kwargs):
+        """
+        Load a checkpoint into a new Trainer object and return that
+        object. If the path specified points at a folder containing
+        multiple checkpoints, the latest one will be used.
+        The dataset can not be serialized and saved so it is required
+        to be explicitly given when loading a checkpoint.
+        Arguments:
+            checkpoint_path (str): Path to a checkpoint or to a folder
+                containing one or more checkpoints.
+            dataset (indexable): The dataset to use.
+            **kwargs (keyword arguments): Any other arguments to override
+                the ones saved in the checkpoint. Useful for when training
+                is continued on a different device or when distributed training
+                is changed.
+        """
+        checkpoint_path = _find_checkpoint(checkpoint_path)
+        _is_checkpoint(checkpoint_path, enforce=True)
+        with open(os.path.join(checkpoint_path, 'kwargs.json'), 'r') as fp:
+            loaded_kwargs = json.load(fp)
+        loaded_kwargs.update(**kwargs)
+        device = torch.device('cpu')
+        if isinstance(loaded_kwargs['device'], (list, tuple)):
+            device = torch.device(loaded_kwargs['device'][0])
+        for name in ['G', 'D']:
+            fpath = os.path.join(checkpoint_path, name + '.pth')
+            loaded_kwargs[name] = models.load(fpath, map_location=device)
+        if os.path.exists(os.path.join(checkpoint_path, 'Gs.pth')):
+            loaded_kwargs['Gs'] = models.load(
+                os.path.join(checkpoint_path, 'Gs.pth'),
+                map_location=device if loaded_kwargs['Gs_device'] is None \
+                    else torch.device(loaded_kwargs['Gs_device'])
+            )
+        obj = cls(dataset=dataset, **loaded_kwargs)
+        for name in ['G_opt', 'D_opt']:
+            fpath = os.path.join(checkpoint_path, name + '.pth')
+            state_dict = torch.load(fpath, map_location=device)
+            getattr(obj, name).load_state_dict(state_dict)
+        return obj
+
+
+#----------------------------------------------------------------------------
+# Checkpoint helper functions
+
+
+def _is_checkpoint(dir_path, enforce=False):
+    if not dir_path:
+        if enforce:
+            raise ValueError('Not a checkpoint.')
+        return False
+    if not os.path.exists(dir_path):
+        if enforce:
+            raise FileNotFoundError('{} could not be found.'.format(dir_path))
+        return False
+    if not os.path.isdir(dir_path):
+        if enforce:
+            raise NotADirectoryError('{} is not a directory.'.format(dir_path))
+        return False
+    fnames = os.listdir(dir_path)
+    for fname in ['G.pth', 'D.pth', 'G_opt.pth', 'D_opt.pth', 'kwargs.json']:
+        if fname not in fnames:
+            if enforce:
+                raise FileNotFoundError(
+                    'Could not find {} in {}.'.format(fname, dir_path))
+            return False
+    return True
+
+
+def _find_checkpoint(dir_path):
+    if not dir_path:
+        return None
+    if not os.path.exists(dir_path) or not os.path.isdir(dir_path):
+        return None
+    if _is_checkpoint(dir_path):
+        return dir_path
+    checkpoint_names = []
+    for name in os.listdir(dir_path):
+        if _is_checkpoint(os.path.join(dir_path, name)):
+            checkpoint_names.append(name)
+    if not checkpoint_names:
+        return None
+    def get_iteration(name):
+        return int(name.split('_')[0])
+    def get_timestamp(name):
+        return '_'.join(name.split('_')[1:])
+    # Python sort is stable, meaning that this sort operation
+    # will guarantee that the order of values after the first
+    # sort will stay for a set of values that have the same
+    # key value.
+    checkpoint_names = sorted(
+        sorted(checkpoint_names, key=get_iteration), key=get_timestamp)
+    return os.path.join(dir_path, checkpoint_names[-1])
+
+
+#----------------------------------------------------------------------------
+# Reg and loss function fetchers
+
+
+def build_opt(net, opt_class, opt_kwargs, reg, reg_interval):
+    opt_kwargs['lr'] = opt_kwargs.get('lr', 1e-3)
+    if reg not in [None, False] and reg_interval:
+        mb_ratio = reg_interval / (reg_interval + 1.)
+        opt_kwargs['lr'] *= mb_ratio
+        if 'momentum' in opt_kwargs:
+            opt_kwargs['momentum'] = opt_kwargs['momentum'] ** mb_ratio
+        if 'betas' in opt_kwargs:
+            betas = opt_kwargs['betas']
+            opt_kwargs['betas'] = (betas[0] ** mb_ratio, betas[1] ** mb_ratio)
+    if isinstance(opt_class, str):
+        opt_class = getattr(torch.optim, opt_class.title())
+    return opt_class(net.parameters(), **opt_kwargs)
+
+
+#----------------------------------------------------------------------------
+# Reg and loss function fetchers
+
+
+_LOSS_FNS = {
+    'G': {
+        'logistic': loss_fns.G_logistic,
+        'logistic_ns': loss_fns.G_logistic_ns,
+        'wgan': loss_fns.G_wgan
+    },
+    'D': {
+        'logistic': loss_fns.D_logistic,
+        'wgan': loss_fns.D_wgan
+    }
+}
+def get_loss_fn(net, loss):
+    if callable(loss):
+        return loss
+    net = net.upper()
+    assert net in ['G', 'D'], 'Unknown net type {}'.format(net)
+    loss = loss.lower()
+    for name in _LOSS_FNS[net].keys():
+        if loss == name:
+            return _LOSS_FNS[net][name]
+    raise ValueError('Unknow {} loss {}'.format(net, loss))
+
+
+_REG_FNS = {
+    'G': {
+        'pathreg': loss_fns.G_pathreg
+    },
+    'D': {
+        'r1': loss_fns.D_r1,
+        'r2': loss_fns.D_r2,
+        'gp': loss_fns.D_gp,
+    }
+}
+def get_reg_fn(net, reg, **kwargs):
+    if reg is None:
+        return None
+    if callable(reg):
+        functools.partial(reg, **kwargs)
+    net = net.upper()
+    assert net in ['G', 'D'], 'Unknown net type {}'.format(net)
+    reg = reg.lower()
+    gamma = None
+    for name in _REG_FNS[net].keys():
+        if reg.startswith(name):
+            gamma_chars = [c for c in reg.replace(name, '') if c.isdigit() or c == '.']
+            if gamma_chars:
+                kwargs.update(gamma=float(''.join(gamma_chars)))
+            return functools.partial(_REG_FNS[net][name], **kwargs)
+    raise ValueError('Unknow regularizer {}'.format(reg))

stylegan2/utils.py

@@ -0,0 +1,726 @@
+import time
+import numbers
+import re
+import sys
+import collections
+import argparse
+import yaml
+from PIL import Image
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+import torchvision
+try:
+    import tqdm
+except ImportError:
+    pass
+try:
+    from IPython.display import display as notebook_display
+    from IPython.display import clear_output as notebook_clear
+except ImportError:
+    pass
+
+
+#----------------------------------------------------------------------------
+# Miscellaneous utils
+
+
+class AttributeDict(dict):
+    """
+    Dict where values can be accessed using attribute syntax.
+    Same as "EasyDict" in the NVIDIA stylegan git repository.
+    """
+
+    def __getattr__(self, name):
+        try:
+            return self[name]
+        except KeyError:
+            raise AttributeError(name)
+
+    def __setattr__(self, name, value):
+        self[name] = value
+
+    def __delattr__(self, name):
+        del self[name]
+
+    def __getstate__(self):
+        return dict(**self)
+
+    def __setstate__(self, state):
+        self.update(**state)
+
+    def __repr__(self):
+        return '{}({})'.format(
+            self.__class__.__name__,
+            ', '.join('{}={}'.format(key, value) for key, value in self.items())
+        )
+
+    @classmethod
+    def convert_dict_recursive(cls, obj):
+        if isinstance(obj, dict):
+            for key in list(obj.keys()):
+                obj[key] = cls.convert_dict_recursive(obj[key])
+            if not isinstance(obj, cls):
+                return cls(**obj)
+        return obj
+
+
+class Timer:
+
+    def __init__(self):
+        self.reset()
+
+    def __enter__(self):
+        self._t0 = time.time()
+
+    def __exit__(self, *args):
+        self._t += time.time() - self._t0
+
+    def value(self):
+        return self._t
+
+    def reset(self):
+        self._t = 0
+
+    def __str__(self):
+        """
+        Get a string representation of the recorded time.
+        Returns:
+            time_as_string (str)
+        """
+        value = self.value()
+        if not value or value >= 100:
+            return '{} s'.format(int(value))
+        elif value >= 1:
+            return '{:.3g} s'.format(value)
+        elif value >= 1e-3:
+            return '{:.3g} ms'.format(value * 1e+3)
+        elif value >= 1e-6:
+            return '{:.3g} us'.format(value * 1e+6)
+        elif value >= 1e-9:
+            return '{:.3g} ns'.format(value * 1e+9)
+        else:
+            return '{:.2E} s'.format(value)
+
+
+def to_list(values):
+    if values is None:
+        return []
+    if isinstance(values, tuple):
+        return list(values)
+    if not isinstance(values, list):
+        return [values]
+    return values
+
+
+def lerp(a, b, beta):
+    if isinstance(beta, numbers.Number):
+        if beta == 1:
+            return b
+        elif beta == 0:
+            return a
+    if torch.is_tensor(a) and a.dtype == torch.float32:
+        # torch lerp only available for fp32
+        return torch.lerp(a, b, beta)
+    # More numerically stable than a + beta * (b - a)
+    return (1 - beta) * a + beta * b
+
+
+def _normalize(v):
+    return v * torch.rsqrt(torch.sum(v ** 2, dim=-1, keepdim=True))
+
+
+def slerp(a, b, beta):
+    assert a.size() == b.size(), 'Size mismatch between ' + \
+        'slerp arguments, received {} and {}'.format(a.size(), b.size())
+    if not torch.is_tensor(beta):
+        beta = torch.tensor(beta).to(a)
+    a = _normalize(a)
+    b = _normalize(b)
+    d = torch.sum(a * b, axis=-1, keepdim=True)
+    p = beta * torch.acos(beta)
+    c = _normalize(b - d * a)
+    d = a * torch.cos(p) + c * torch.sin(p)
+    return _normalize(d)
+
+
+#----------------------------------------------------------------------------
+# Command line utils
+
+
+def _parse_configs(configs):
+    kwargs = {}
+    for config in configs:
+        with open(config, 'r') as fp:
+            kwargs.update(yaml.safe_load(fp))
+    return kwargs
+
+
+class ConfigArgumentParser(argparse.ArgumentParser):
+
+    _CONFIG_ARG_KEY = '_configs'
+
+    def __init__(self, *args, **kwargs):
+        super(ConfigArgumentParser, self).__init__(*args, **kwargs)
+        self.add_argument(
+            self._CONFIG_ARG_KEY,
+            nargs='*',
+            help='Any yaml-style config file whos values will override the defaults of this argument parser.',
+            type=str
+        )
+
+    def parse_args(self, args=None):
+        config_args = _parse_configs(
+            getattr(
+                super(ConfigArgumentParser, self).parse_args(args),
+                self._CONFIG_ARG_KEY
+            )
+        )
+        self.set_defaults(**config_args)
+        return super(ConfigArgumentParser, self).parse_args(args)
+
+
+def bool_type(v):
+    if isinstance(v, bool):
+       return v
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+
+def range_type(s):
+    """
+    Accept either a comma separated list of numbers
+    'a,b,c' or a range 'a-c' and return as a list of ints.
+    """
+    range_re = re.compile(r'^(\d+)-(\d+)$')
+    m = range_re.match(s)
+    if m:
+        return range(int(m.group(1)), int(m.group(2))+1)
+    vals = s.split(',')
+    return [int(x) for x in vals]
+
+
+#----------------------------------------------------------------------------
+# Dataset and generation of latents
+
+
+class ResizeTransform:
+
+    def __init__(self, height, width, resize=True, mode='bicubic'):
+        if resize:
+            assert height and width, 'Height and width have to be given ' + \
+                'when resizing data.'
+        self.height = height
+        self.width = width
+        self.resize = resize
+        self.mode = mode
+
+    def __call__(self, tensor):
+        if self.height and self.width:
+            if tensor.size(1) != self.height or tensor.size(2) != self.width:
+                if self.resize:
+                    kwargs = {}
+                    if 'cubic' in self.mode or 'linear' in self.mode:
+                        kwargs.update(align_corners=False)
+                    tensor = F.interpolate(
+                        tensor.unsqueeze(0),
+                        size=(self.height, self.width),
+                        mode=self.mode,
+                        **kwargs
+                    ).squeeze(0)
+                else:
+                    raise ValueError(
+                        'Data shape incorrect, expected ({},{}) '.format(self.width, self.height) + \
+                        'but got ({},{}) (width, height)'.format(tensor.size(2), tensor.size(1))
+                    )
+        return tensor
+
+
+def _PIL_RGB_loader(path):
+    return Image.open(path).convert('RGB')
+
+
+def _PIL_grayscale_loader(path):
+    return Image.open(path).convert('L')
+
+
+class ImageFolder(torchvision.datasets.ImageFolder):
+
+    def __init__(self,
+                 *args,
+                 mirror=False,
+                 pixel_min=-1,
+                 pixel_max=1,
+                 height=None,
+                 width=None,
+                 resize=False,
+                 resize_mode='bicubic',
+                 grayscale=False,
+                 **kwargs):
+        super(ImageFolder, self).__init__(
+            *args,
+            loader=_PIL_grayscale_loader if grayscale else _PIL_RGB_loader,
+            **kwargs
+        )
+        transforms = []
+        if mirror:
+            transforms.append(torchvision.transforms.RandomHorizontalFlip())
+        transforms.append(torchvision.transforms.ToTensor())
+        transforms.append(
+            torchvision.transforms.Normalize(
+                mean=[-(pixel_min / (pixel_max - pixel_min))],
+                std=[1. / (pixel_max - pixel_min)]
+            )
+        )
+        transforms.append(ResizeTransform(
+            height=height, width=width, resize=resize, mode=resize_mode))
+        self.transform = torchvision.transforms.Compose(transforms)
+
+    def _find_classes(self, *args, **kwargs):
+        classes, class_to_idx = super(ImageFolder, self)._find_classes(*args, **kwargs)
+        if not classes:
+            classes = ['']
+            class_to_idx = {'': 0}
+        return classes, class_to_idx
+
+
+class PriorGenerator:
+
+    def __init__(self, latent_size, label_size, batch_size, device):
+        self.latent_size = latent_size
+        self.label_size = label_size
+        self.batch_size = batch_size
+        self.device = device
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        return self()
+
+    def __call__(self, batch_size=None, multi_latent_prob=0, seed=None):
+        if batch_size is None:
+            batch_size = self.batch_size
+        shape = [batch_size, self.latent_size]
+        if multi_latent_prob:
+            if seed is not None:
+                np.random.seed(seed)
+            if np.random.uniform() < multi_latent_prob:
+                shape = [batch_size, 2, self.latent_size]
+        if seed is not None:
+            torch.manual_seed(seed)
+        latents = torch.empty(*shape, device=self.device).normal_()
+        labels = None
+        if self.label_size:
+            label_shape = [batch_size]
+            labels = torch.randint(0, self.label_size, label_shape, device=self.device)
+        return latents, labels
+
+
+#----------------------------------------------------------------------------
+# Training utils
+
+
+class MovingAverageModule:
+
+    def __init__(self,
+                 from_module,
+                 to_module=None,
+                 param_beta=0.995,
+                 buffer_beta=0,
+                 device=None):
+        from_module = unwrap_module(from_module)
+        to_module = unwrap_module(to_module)
+        if device is None:
+            module = from_module
+            if to_module is not None:
+                module = to_module
+            device = next(module.parameters()).device
+        else:
+            device = torch.device(device)
+        self.from_module = from_module
+        if to_module is None:
+            self.module = from_module.clone().to(device)
+        else:
+            assert type(to_module) == type(from_module), \
+                'Mismatch between type of source and target module.'
+            assert set(self._get_named_parameters(to_module).keys()) \
+            == set(self._get_named_parameters(from_module).keys()), \
+                'Mismatch between parameters of source and target module.'
+            assert set(self._get_named_buffers(to_module).keys()) \
+            == set(self._get_named_buffers(from_module).keys()), \
+                'Mismatch between buffers of source and target module.'
+            self.module = to_module.to(device)
+        self.module.eval().requires_grad_(False)
+        self.param_beta = param_beta
+        self.buffer_beta = buffer_beta
+        self.device = device
+
+    def __getattr__(self, name):
+        try:
+            return super(object, self).__getattr__(name)
+        except AttributeError:
+            return getattr(self.module, name)
+
+    def update(self):
+        self._update_data(
+            from_data=self._get_named_parameters(self.from_module),
+            to_data=self._get_named_parameters(self.module),
+            beta=self.param_beta
+        )
+        self._update_data(
+            from_data=self._get_named_buffers(self.from_module),
+            to_data=self._get_named_buffers(self.module),
+            beta=self.buffer_beta
+        )
+
+    @staticmethod
+    def _update_data(from_data, to_data, beta):
+        for name in from_data.keys():
+            if name not in to_data:
+                continue
+            fr, to = from_data[name], to_data[name]
+            with torch.no_grad():
+                if beta == 0:
+                    to.data.copy_(fr.data.to(to.data))
+                elif beta < 1:
+                    to.data.copy_(lerp(fr.data.to(to.data), to.data, beta))
+
+    @staticmethod
+    def _get_named_parameters(module):
+        return {name: value for name, value in module.named_parameters()}
+
+    @staticmethod
+    def _get_named_buffers(module):
+        return {name: value for name, value in module.named_buffers()}
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
+
+    def forward(self, *args, **kwargs):
+        self.module.eval()
+        args, args_in_device = move_to_device(args, self.device)
+        kwargs, kwargs_in_device = move_to_device(kwargs, self.device)
+        in_device = None
+        if args_in_device is not None:
+            in_device = args_in_device
+        if kwargs_in_device is not None:
+            in_device = kwargs_in_device
+        out = self.module(*args, **kwargs)
+        if in_device is not None:
+            out, _ = move_to_device(out, in_device)
+        return out
+
+
+def move_to_device(value, device):
+    if torch.is_tensor(value):
+        value.to(device), value.device
+    orig_device = None
+    if isinstance(value, (tuple, list)):
+        values = []
+        for val in value:
+            _val, orig_device = move_to_device(val, device)
+            values.append(_val)
+        return type(value)(values), orig_device
+    if isinstance(value, dict):
+        if isinstance(value, collections.OrderedDict):
+            values = collections.OrderedDict()
+        else:
+            values = {}
+        for key, val in value.items():
+            _val, orig_device = move_to_device(val, device)
+            values[key] = val
+        return values, orig_device
+    return value, orig_device
+
+
+_WRAPPER_CLASSES = (MovingAverageModule, nn.DataParallel, nn.parallel.DistributedDataParallel)
+def unwrap_module(module):
+    if isinstance(module, _WRAPPER_CLASSES):
+        return module.module
+    return module
+
+
+def get_grad_norm_from_optimizer(optimizer, norm_type=2):
+    """
+    Get the gradient norm for some parameters contained in an optimizer.
+    Arguments:
+        optimizer (torch.optim.Optimizer)
+        norm_type (int): Type of norm. Default value is 2.
+    Returns:
+        norm (float)
+    """
+    total_norm = 0
+    if optimizer is not None:
+        for param_group in optimizer.param_groups:
+            for p in param_group['params']:
+                if p.grad is not None:
+                    with torch.no_grad():
+                        param_norm = p.grad.data.norm(norm_type)
+                    total_norm += param_norm ** norm_type
+    total_norm = total_norm ** (1. / norm_type)
+    return total_norm.item()
+
+
+#----------------------------------------------------------------------------
+# printing and logging utils
+
+
+class ValueTracker:
+
+    def __init__(self, beta=0.95):
+        self.beta = beta
+        self.values = {}
+
+    def add(self, name, value, beta=None):
+        if torch.is_tensor(value):
+            value = value.item()
+        if beta is None:
+            beta = self.beta
+        if name not in self.values:
+            self.values[name] = value
+        else:
+            self.values[name] = lerp(value, self.values[name], beta)
+
+    def __getitem__(self, key):
+        return self.values[key]
+
+    def __str__(self):
+        string = ''
+        for i, name in enumerate(sorted(self.values.keys())):
+            if i and i % 3 == 0:
+                string += '\n'
+            elif string:
+                string += ', '
+            format_string = '{}: {}'
+            if isinstance(self.values[name], float):
+                format_string = '{}: {:.4g}'
+            string += format_string.format(name, self.values[name])
+        return string
+
+
+def is_notebook():
+    """
+    Check if code is running from jupyter notebook.
+    Returns:
+        notebook (bool): True if running from jupyter notebook,
+            else False.
+    """
+    try:
+        __IPYTHON__
+        return True
+    except NameError:
+        return False
+
+
+def _progress_bar(count, total):
+    """
+    Get a simple one-line string representing a progress bar.
+    Arguments:
+        count (int): Current count. Starts at 0.
+        total (int): Total count.
+    Returns:
+        pbar_string (str): The string progress bar.
+    """
+    bar_len = 60
+    filled_len = int(round(bar_len * (count + 1) / float(total)))
+    bar = '=' * filled_len + '-' * (bar_len - filled_len)
+    return '[{}] {}/{}'.format(bar, count + 1, total)
+
+
+class ProgressWriter:
+    """
+    Handles writing output and displaying a progress bar. Automatically
+    adjust for notebooks. Supports outputting text
+    that is compatible with the progressbar (in notebooks the text is
+    refreshed instead of printed).
+    Arguments:
+        length (int, optional): Total length of the progressbar.
+            Default value is None.
+        progress_bar (bool, optional): Display a progressbar.
+            Default value is True.
+        clear (bool, optional): If running from a notebook, clear
+            the current cell's output. Default value is False.
+    """
+    def __init__(self, length=None, progress_bar=True, clear=False):
+        if is_notebook() and clear:
+            notebook_clear()
+
+        if length is not None:
+            length = int(length)
+        self.length = length
+        self.count = 0
+
+        self._simple_pbar = False
+        if progress_bar and 'tqdm' not in sys.modules:
+            self._simple_pbar = True
+
+        progress_bar = progress_bar and 'tqdm' in sys.modules
+
+        self._progress_bar = None
+        if progress_bar:
+            pbar = tqdm.tqdm
+            if is_notebook():
+                pbar = tqdm.tqdm_notebook
+            if length is not None:
+                self._progress_bar = pbar(total=length, file=sys.stdout)
+            else:
+                self._progress_bar = pbar(file=sys.stdout)
+
+        if is_notebook():
+            self._writer = notebook_display(
+                _StrRepr(''),
+                display_id=time.asctime()
+            )
+        else:
+            if progress_bar:
+                self._writer = self._progress_bar
+            else:
+                self._writer = sys.stdout
+
+    def write(self, *lines, step=True):
+        """
+        Output values to stdout (or a display object if called from notebook).
+        Arguments:
+            *lines: The lines to write (positional arguments).
+            step (bool): Update the progressbar if present.
+                Default value is True.
+        """
+        string = '\n'.join(str(line) for line in lines if line and line.strip())
+        if self._simple_pbar:
+            string = _progress_bar(self.count, self.length) + '\n' + string
+        if is_notebook():
+            self._writer.update(_StrRepr(string))
+        else:
+            self._writer.write('\n\n' + string)
+            if hasattr(self._writer, 'flush'):
+                self._writer.flush()
+        if step:
+            self.step()
+
+    def step(self):
+        """
+        Update the progressbar if present.
+        """
+        self.count += 1
+        if self._progress_bar is not None:
+            self._progress_bar.update()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        return next(self.rnge)
+
+    def close(self):
+        if hasattr(self._writer, 'close'):
+            can_close = True
+            try:
+                can_close = self._writer != sys.stdout and self._writer != sys.stderr
+            except AttributeError:
+                pass
+            if can_close:
+                self._writer.close()
+        if hasattr(self._progress_bar, 'close'):
+            self._progress_bar.close()
+
+    def __del__(self):
+        self.close()
+
+
+class _StrRepr:
+    """
+    A wrapper for strings that returns the string
+    on repr() calls. Used by notebooks.
+    """
+    def __init__(self, string):
+        self.string = string
+
+    def __repr__(self):
+        return self.string
+
+
+#----------------------------------------------------------------------------
+# image utils
+
+
+def tensor_to_PIL(image_tensor, pixel_min=-1, pixel_max=1):
+    image_tensor = image_tensor.cpu()
+    if pixel_min != 0 or pixel_max != 1:
+        image_tensor = (image_tensor - pixel_min) / (pixel_max - pixel_min)
+    image_tensor.clamp_(min=0, max=1)
+    to_pil = torchvision.transforms.functional.to_pil_image
+    if image_tensor.dim() == 4:
+        return [to_pil(img) for img in image_tensor]
+    return to_pil(image_tensor)
+
+
+def PIL_to_tensor(image, pixel_min=-1, pixel_max=1):
+    to_tensor = torchvision.transforms.functional.to_tensor
+    if isinstance(image, (list, tuple)):
+        image_tensor = torch.stack([to_tensor(img) for img in image])
+    else:
+        image_tensor = to_tensor(image)
+    if pixel_min != 0 or pixel_max != 1:
+        image_tensor = image_tensor * (pixel_max - pixel_min) + pixel_min
+    return image_tensor
+
+
+def stack_images_PIL(imgs, shape=None, individual_img_size=None):
+    """
+    Concatenate multiple images into a grid within a single image.
+    Arguments:
+        imgs (Sequence of PIL.Image): Input images.
+        shape (list, tuple, int, optional): Shape of the grid. Should consist
+            of two values, (width, height). If an integer value is passed it
+            is used for both width and height. If no value is passed the shape
+            is infered from the number of images. Default value is None.
+        individual_img_size (list, tuple, int, optional): The size of the
+            images being concatenated. Default value is None.
+    Returns:
+        canvas (PIL.Image): Image containing input images in a grid.
+    """
+    assert len(imgs) > 0, 'No images received.'
+    if shape is None:
+        size = int(np.ceil(np.sqrt(len(imgs))))
+        shape = [int(np.ceil(len(imgs) / size)), size]
+    else:
+        if isinstance(shape, numbers.Number):
+            shape = 2 * [shape]
+        assert len(shape) == 2, 'Shape should specify (width, height).'
+
+    if individual_img_size is None:
+        for i in range(len(imgs) - 1):
+            assert imgs[i].size == imgs[i + 1].size, \
+                'Images are of different sizes, please specify a ' + \
+                'size (width, height). Found sizes:\n' + \
+                ', '.join(str(img.size) for img in imgs)
+        individual_img_size = imgs[0].size
+    else:
+        if not isinstance(individual_img_size, (tuple, list)):
+            individual_img_size = 2 * (individual_img_size,)
+        individual_img_size = tuple(individual_img_size)
+        for i in range(len(imgs)):
+            if imgs[i].size != individual_img_size:
+                imgs[i] = imgs[i].resize(individual_img_size)
+
+    width, height = individual_img_size
+    width, height = int(width), int(height)
+    canvas = Image.new(
+        'RGB',
+        (shape[0] * width, shape[1] * height),
+        (0, 0, 0, 0)
+    )
+    imgs = imgs.copy()
+    for h_i in range(shape[1]):
+        for w_i in range(shape[0]):
+            if len(imgs) > 0:
+                img = imgs.pop(0).convert('RGB')
+                offset = (w_i * width, h_i * height)
+                canvas.paste(img, offset)
+    return canvas