Added base files

align_images.py

@@ -0,0 +1,57 @@
+import os
+import sys
+import bz2
+import argparse
+from keras.utils import get_file
+from ffhq_dataset.face_alignment import image_align
+from ffhq_dataset.landmarks_detector import LandmarksDetector
+import multiprocessing
+
+def unpack_bz2(src_path):
+    data = bz2.BZ2File(src_path).read()
+    dst_path = src_path[:-4]
+    with open(dst_path, 'wb') as fp:
+        fp.write(data)
+    return dst_path
+
+
+if __name__ == "__main__":
+    """
+    Extracts and aligns all faces from images using DLib and a function from original FFHQ dataset preparation step
+    python align_images.py /raw_images /aligned_images
+    """
+    parser = argparse.ArgumentParser(description='Align faces from input images', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument('raw_dir', help='Directory with raw images for face alignment')
+    parser.add_argument('aligned_dir', help='Directory for storing aligned images')
+    parser.add_argument('--output_size', default=1024, help='The dimension of images for input to the model', type=int)
+    parser.add_argument('--x_scale', default=1, help='Scaling factor for x dimension', type=float)
+    parser.add_argument('--y_scale', default=1, help='Scaling factor for y dimension', type=float)
+    parser.add_argument('--em_scale', default=0.1, help='Scaling factor for eye-mouth distance', type=float)
+    parser.add_argument('--use_alpha', default=False, help='Add an alpha channel for masking', type=bool)
+
+    args, other_args = parser.parse_known_args()
+
+    landmarks_model_path = unpack_bz2("shape_predictor_68_face_landmarks.dat.bz2")
+    RAW_IMAGES_DIR = args.raw_dir
+    ALIGNED_IMAGES_DIR = args.aligned_dir
+
+    landmarks_detector = LandmarksDetector(landmarks_model_path)
+    for img_name in os.listdir(RAW_IMAGES_DIR):
+        print('Aligning %s ...' % img_name)
+        try:
+            raw_img_path = os.path.join(RAW_IMAGES_DIR, img_name)
+            fn = face_img_name = '%s_%02d.png' % (os.path.splitext(img_name)[0], 1)
+            if os.path.isfile(fn):
+                continue
+            print('Getting landmarks...')
+            for i, face_landmarks in enumerate(landmarks_detector.get_landmarks(raw_img_path), start=1):
+                try:
+                    print('Starting face alignment...')
+                    face_img_name = '%s_%02d.png' % (os.path.splitext(img_name)[0], i)
+                    aligned_face_path = os.path.join(ALIGNED_IMAGES_DIR, face_img_name)
+                    image_align(raw_img_path, aligned_face_path, face_landmarks, output_size=args.output_size, x_scale=args.x_scale, y_scale=args.y_scale, em_scale=args.em_scale, alpha=args.use_alpha)
+                    print('Wrote result %s' % aligned_face_path)
+                except:
+                    print("Exception in face alignment!")
+        except:
+            print("Exception in landmark detection!")

config.py

@@ -0,0 +1,22 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial
+# 4.0 International License. To view a copy of this license, visit
+# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
+# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
+
+"""Global configuration."""
+
+#----------------------------------------------------------------------------
+# Paths.
+
+result_dir = 'results'
+data_dir = 'datasets'
+cache_dir = 'cache'
+run_dir_ignore = ['results', 'datasets', 'cache']
+
+# experimental - replace Dense layers with TreeConnect
+use_treeconnect = False
+treeconnect_threshold = 1024
+
+#----------------------------------------------------------------------------

dnnlib/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial
+# 4.0 International License. To view a copy of this license, visit
+# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
+# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
+
+from . import submission
+
+from .submission.run_context import RunContext
+
+from .submission.submit import SubmitTarget
+from .submission.submit import PathType
+from .submission.submit import SubmitConfig
+from .submission.submit import get_path_from_template
+from .submission.submit import submit_run
+
+from .util import EasyDict
+
+submit_config: SubmitConfig = None # Package level variable for SubmitConfig which is only valid when inside the run function.

dnnlib/submission/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial
+# 4.0 International License. To view a copy of this license, visit
+# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
+# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
+
+from . import run_context
+from . import submit

dnnlib/submission/_internal/run.py

@@ -0,0 +1,45 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial
+# 4.0 International License. To view a copy of this license, visit
+# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
+# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
+
+"""Helper for launching run functions in computing clusters.
+
+During the submit process, this file is copied to the appropriate run dir.
+When the job is launched in the cluster, this module is the first thing that
+is run inside the docker container.
+"""
+
+import os
+import pickle
+import sys
+
+# PYTHONPATH should have been set so that the run_dir/src is in it
+import dnnlib
+
+def main():
+    if not len(sys.argv) >= 4:
+        raise RuntimeError("This script needs three arguments: run_dir, task_name and host_name!")
+
+    run_dir = str(sys.argv[1])
+    task_name = str(sys.argv[2])
+    host_name = str(sys.argv[3])
+
+    submit_config_path = os.path.join(run_dir, "submit_config.pkl")
+
+    # SubmitConfig should have been pickled to the run dir
+    if not os.path.exists(submit_config_path):
+        raise RuntimeError("SubmitConfig pickle file does not exist!")
+
+    submit_config: dnnlib.SubmitConfig = pickle.load(open(submit_config_path, "rb"))
+    dnnlib.submission.submit.set_user_name_override(submit_config.user_name)
+
+    submit_config.task_name = task_name
+    submit_config.host_name = host_name
+
+    dnnlib.submission.submit.run_wrapper(submit_config)
+
+if __name__ == "__main__":
+    main()

dnnlib/submission/run_context.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial
+# 4.0 International License. To view a copy of this license, visit
+# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
+# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
+
+"""Helpers for managing the run/training loop."""
+
+import datetime
+import json
+import os
+import pprint
+import time
+import types
+
+from typing import Any
+
+from . import submit
+
+
+class RunContext(object):
+    """Helper class for managing the run/training loop.
+
+    The context will hide the implementation details of a basic run/training loop.
+    It will set things up properly, tell if run should be stopped, and then cleans up.
+    User should call update periodically and use should_stop to determine if run should be stopped.
+
+    Args:
+        submit_config: The SubmitConfig that is used for the current run.
+        config_module: The whole config module that is used for the current run.
+        max_epoch: Optional cached value for the max_epoch variable used in update.
+    """
+
+    def __init__(self, submit_config: submit.SubmitConfig, config_module: types.ModuleType = None, max_epoch: Any = None):
+        self.submit_config = submit_config
+        self.should_stop_flag = False
+        self.has_closed = False
+        self.start_time = time.time()
+        self.last_update_time = time.time()
+        self.last_update_interval = 0.0
+        self.max_epoch = max_epoch
+
+        # pretty print the all the relevant content of the config module to a text file
+        if config_module is not None:
+            with open(os.path.join(submit_config.run_dir, "config.txt"), "w") as f:
+                filtered_dict = {k: v for k, v in config_module.__dict__.items() if not k.startswith("_") and not isinstance(v, (types.ModuleType, types.FunctionType, types.LambdaType, submit.SubmitConfig, type))}
+                pprint.pprint(filtered_dict, stream=f, indent=4, width=200, compact=False)
+
+        # write out details about the run to a text file
+        self.run_txt_data = {"task_name": submit_config.task_name, "host_name": submit_config.host_name, "start_time": datetime.datetime.now().isoformat(sep=" ")}
+        with open(os.path.join(submit_config.run_dir, "run.txt"), "w") as f:
+            pprint.pprint(self.run_txt_data, stream=f, indent=4, width=200, compact=False)
+
+    def __enter__(self) -> "RunContext":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.close()
+
+    def update(self, loss: Any = 0, cur_epoch: Any = 0, max_epoch: Any = None) -> None:
+        """Do general housekeeping and keep the state of the context up-to-date.
+        Should be called often enough but not in a tight loop."""
+        assert not self.has_closed
+
+        self.last_update_interval = time.time() - self.last_update_time
+        self.last_update_time = time.time()
+
+        if os.path.exists(os.path.join(self.submit_config.run_dir, "abort.txt")):
+            self.should_stop_flag = True
+
+        max_epoch_val = self.max_epoch if max_epoch is None else max_epoch
+
+    def should_stop(self) -> bool:
+        """Tell whether a stopping condition has been triggered one way or another."""
+        return self.should_stop_flag
+
+    def get_time_since_start(self) -> float:
+        """How much time has passed since the creation of the context."""
+        return time.time() - self.start_time
+
+    def get_time_since_last_update(self) -> float:
+        """How much time has passed since the last call to update."""
+        return time.time() - self.last_update_time
+
+    def get_last_update_interval(self) -> float:
+        """How much time passed between the previous two calls to update."""
+        return self.last_update_interval
+
+    def close(self) -> None:
+        """Close the context and clean up.
+        Should only be called once."""
+        if not self.has_closed:
+            # update the run.txt with stopping time
+            self.run_txt_data["stop_time"] = datetime.datetime.now().isoformat(sep=" ")
+            with open(os.path.join(self.submit_config.run_dir, "run.txt"), "w") as f:
+                pprint.pprint(self.run_txt_data, stream=f, indent=4, width=200, compact=False)
+
+            self.has_closed = True

dnnlib/submission/submit.py

@@ -0,0 +1,290 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial
+# 4.0 International License. To view a copy of this license, visit
+# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
+# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
+
+"""Submit a function to be run either locally or in a computing cluster."""
+
+import copy
+import io
+import os
+import pathlib
+import pickle
+import platform
+import pprint
+import re
+import shutil
+import time
+import traceback
+
+import zipfile
+
+from enum import Enum
+
+from .. import util
+from ..util import EasyDict
+
+
+class SubmitTarget(Enum):
+    """The target where the function should be run.
+
+    LOCAL: Run it locally.
+    """
+    LOCAL = 1
+
+
+class PathType(Enum):
+    """Determines in which format should a path be formatted.
+
+    WINDOWS: Format with Windows style.
+    LINUX: Format with Linux/Posix style.
+    AUTO: Use current OS type to select either WINDOWS or LINUX.
+    """
+    WINDOWS = 1
+    LINUX = 2
+    AUTO = 3
+
+
+_user_name_override = None
+
+
+class SubmitConfig(util.EasyDict):
+    """Strongly typed config dict needed to submit runs.
+
+    Attributes:
+        run_dir_root: Path to the run dir root. Can be optionally templated with tags. Needs to always be run through get_path_from_template.
+        run_desc: Description of the run. Will be used in the run dir and task name.
+        run_dir_ignore: List of file patterns used to ignore files when copying files to the run dir.
+        run_dir_extra_files: List of (abs_path, rel_path) tuples of file paths. rel_path root will be the src directory inside the run dir.
+        submit_target: Submit target enum value. Used to select where the run is actually launched.
+        num_gpus: Number of GPUs used/requested for the run.
+        print_info: Whether to print debug information when submitting.
+        ask_confirmation: Whether to ask a confirmation before submitting.
+        run_id: Automatically populated value during submit.
+        run_name: Automatically populated value during submit.
+        run_dir: Automatically populated value during submit.
+        run_func_name: Automatically populated value during submit.
+        run_func_kwargs: Automatically populated value during submit.
+        user_name: Automatically populated value during submit. Can be set by the user which will then override the automatic value.
+        task_name: Automatically populated value during submit.
+        host_name: Automatically populated value during submit.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+        # run (set these)
+        self.run_dir_root = ""  # should always be passed through get_path_from_template
+        self.run_desc = ""
+        self.run_dir_ignore = ["__pycache__", "*.pyproj", "*.sln", "*.suo", ".cache", ".idea", ".vs", ".vscode"]
+        self.run_dir_extra_files = None
+
+        # submit (set these)
+        self.submit_target = SubmitTarget.LOCAL
+        self.num_gpus = 1
+        self.print_info = False
+        self.ask_confirmation = False
+
+        # (automatically populated)
+        self.run_id = None
+        self.run_name = None
+        self.run_dir = None
+        self.run_func_name = None
+        self.run_func_kwargs = None
+        self.user_name = None
+        self.task_name = None
+        self.host_name = "localhost"
+
+
+def get_path_from_template(path_template: str, path_type: PathType = PathType.AUTO) -> str:
+    """Replace tags in the given path template and return either Windows or Linux formatted path."""
+    # automatically select path type depending on running OS
+    if path_type == PathType.AUTO:
+        if platform.system() == "Windows":
+            path_type = PathType.WINDOWS
+        elif platform.system() == "Linux":
+            path_type = PathType.LINUX
+        else:
+            raise RuntimeError("Unknown platform")
+
+    path_template = path_template.replace("<USERNAME>", get_user_name())
+
+    # return correctly formatted path
+    if path_type == PathType.WINDOWS:
+        return str(pathlib.PureWindowsPath(path_template))
+    elif path_type == PathType.LINUX:
+        return str(pathlib.PurePosixPath(path_template))
+    else:
+        raise RuntimeError("Unknown platform")
+
+
+def get_template_from_path(path: str) -> str:
+    """Convert a normal path back to its template representation."""
+    # replace all path parts with the template tags
+    path = path.replace("\\", "/")
+    return path
+
+
+def convert_path(path: str, path_type: PathType = PathType.AUTO) -> str:
+    """Convert a normal path to template and the convert it back to a normal path with given path type."""
+    path_template = get_template_from_path(path)
+    path = get_path_from_template(path_template, path_type)
+    return path
+
+
+def set_user_name_override(name: str) -> None:
+    """Set the global username override value."""
+    global _user_name_override
+    _user_name_override = name
+
+
+def get_user_name():
+    """Get the current user name."""
+    if _user_name_override is not None:
+        return _user_name_override
+    elif platform.system() == "Windows":
+        return os.getlogin()
+    elif platform.system() == "Linux":
+        try:
+            import pwd # pylint: disable=import-error
+            return pwd.getpwuid(os.geteuid()).pw_name # pylint: disable=no-member
+        except:
+            return "unknown"
+    else:
+        raise RuntimeError("Unknown platform")
+
+
+def _create_run_dir_local(submit_config: SubmitConfig) -> str:
+    """Create a new run dir with increasing ID number at the start."""
+    run_dir_root = get_path_from_template(submit_config.run_dir_root, PathType.AUTO)
+
+    if not os.path.exists(run_dir_root):
+        print("Creating the run dir root: {}".format(run_dir_root))
+        os.makedirs(run_dir_root)
+
+    submit_config.run_id = _get_next_run_id_local(run_dir_root)
+    submit_config.run_name = "{0:05d}-{1}".format(submit_config.run_id, submit_config.run_desc)
+    run_dir = os.path.join(run_dir_root, submit_config.run_name)
+
+    if os.path.exists(run_dir):
+        raise RuntimeError("The run dir already exists! ({0})".format(run_dir))
+
+    print("Creating the run dir: {}".format(run_dir))
+    os.makedirs(run_dir)
+
+    return run_dir
+
+
+def _get_next_run_id_local(run_dir_root: str) -> int:
+    """Reads all directory names in a given directory (non-recursive) and returns the next (increasing) run id. Assumes IDs are numbers at the start of the directory names."""
+    dir_names = [d for d in os.listdir(run_dir_root) if os.path.isdir(os.path.join(run_dir_root, d))]
+    r = re.compile("^\\d+")  # match one or more digits at the start of the string
+    run_id = 0
+
+    for dir_name in dir_names:
+        m = r.match(dir_name)
+
+        if m is not None:
+            i = int(m.group())
+            run_id = max(run_id, i + 1)
+
+    return run_id
+
+
+def _populate_run_dir(run_dir: str, submit_config: SubmitConfig) -> None:
+    """Copy all necessary files into the run dir. Assumes that the dir exists, is local, and is writable."""
+    print("Copying files to the run dir")
+    files = []
+
+    run_func_module_dir_path = util.get_module_dir_by_obj_name(submit_config.run_func_name)
+    assert '.' in submit_config.run_func_name
+    for _idx in range(submit_config.run_func_name.count('.') - 1):
+        run_func_module_dir_path = os.path.dirname(run_func_module_dir_path)
+    files += util.list_dir_recursively_with_ignore(run_func_module_dir_path, ignores=submit_config.run_dir_ignore, add_base_to_relative=False)
+
+    dnnlib_module_dir_path = util.get_module_dir_by_obj_name("dnnlib")
+    files += util.list_dir_recursively_with_ignore(dnnlib_module_dir_path, ignores=submit_config.run_dir_ignore, add_base_to_relative=True)
+
+    if submit_config.run_dir_extra_files is not None:
+        files += submit_config.run_dir_extra_files
+
+    files = [(f[0], os.path.join(run_dir, "src", f[1])) for f in files]
+    files += [(os.path.join(dnnlib_module_dir_path, "submission", "_internal", "run.py"), os.path.join(run_dir, "run.py"))]
+
+    util.copy_files_and_create_dirs(files)
+
+    pickle.dump(submit_config, open(os.path.join(run_dir, "submit_config.pkl"), "wb"))
+
+    with open(os.path.join(run_dir, "submit_config.txt"), "w") as f:
+        pprint.pprint(submit_config, stream=f, indent=4, width=200, compact=False)
+
+
+def run_wrapper(submit_config: SubmitConfig) -> None:
+    """Wrap the actual run function call for handling logging, exceptions, typing, etc."""
+    is_local = submit_config.submit_target == SubmitTarget.LOCAL
+
+    checker = None
+
+    # when running locally, redirect stderr to stdout, log stdout to a file, and force flushing
+    if is_local:
+        logger = util.Logger(file_name=os.path.join(submit_config.run_dir, "log.txt"), file_mode="w", should_flush=True)
+    else:  # when running in a cluster, redirect stderr to stdout, and just force flushing (log writing is handled by run.sh)
+        logger = util.Logger(file_name=None, should_flush=True)
+
+    import dnnlib
+    dnnlib.submit_config = submit_config
+
+    try:
+        print("dnnlib: Running {0}() on {1}...".format(submit_config.run_func_name, submit_config.host_name))
+        start_time = time.time()
+        util.call_func_by_name(func_name=submit_config.run_func_name, submit_config=submit_config, **submit_config.run_func_kwargs)
+        print("dnnlib: Finished {0}() in {1}.".format(submit_config.run_func_name, util.format_time(time.time() - start_time)))
+    except:
+        if is_local:
+            raise
+        else:
+            traceback.print_exc()
+
+            log_src = os.path.join(submit_config.run_dir, "log.txt")
+            log_dst = os.path.join(get_path_from_template(submit_config.run_dir_root), "{0}-error.txt".format(submit_config.run_name))
+            shutil.copyfile(log_src, log_dst)
+    finally:
+        open(os.path.join(submit_config.run_dir, "_finished.txt"), "w").close()
+
+    dnnlib.submit_config = None
+    logger.close()
+
+    if checker is not None:
+        checker.stop()
+
+
+def submit_run(submit_config: SubmitConfig, run_func_name: str, **run_func_kwargs) -> None:
+    """Create a run dir, gather files related to the run, copy files to the run dir, and launch the run in appropriate place."""
+    submit_config = copy.copy(submit_config)
+
+    if submit_config.user_name is None:
+        submit_config.user_name = get_user_name()
+
+    submit_config.run_func_name = run_func_name
+    submit_config.run_func_kwargs = run_func_kwargs
+
+    assert submit_config.submit_target == SubmitTarget.LOCAL
+    if submit_config.submit_target in {SubmitTarget.LOCAL}:
+        run_dir = _create_run_dir_local(submit_config)
+
+        submit_config.task_name = "{0}-{1:05d}-{2}".format(submit_config.user_name, submit_config.run_id, submit_config.run_desc)
+        submit_config.run_dir = run_dir
+        _populate_run_dir(run_dir, submit_config)
+
+    if submit_config.print_info:
+        print("\nSubmit config:\n")
+        pprint.pprint(submit_config, indent=4, width=200, compact=False)
+        print()
+
+    if submit_config.ask_confirmation:
+        if not util.ask_yes_no("Continue submitting the job?"):
+            return
+
+    run_wrapper(submit_config)

dnnlib/tflib/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial
+# 4.0 International License. To view a copy of this license, visit
+# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
+# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
+
+from . import autosummary
+from . import network
+from . import optimizer
+from . import tfutil
+
+from .tfutil import *
+from .network import Network
+
+from .optimizer import Optimizer

dnnlib/tflib/autosummary.py

@@ -0,0 +1,184 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial
+# 4.0 International License. To view a copy of this license, visit
+# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
+# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
+
+"""Helper for adding automatically tracked values to Tensorboard.
+
+Autosummary creates an identity op that internally keeps track of the input
+values and automatically shows up in TensorBoard. The reported value
+represents an average over input components. The average is accumulated
+constantly over time and flushed when save_summaries() is called.
+
+Notes:
+- The output tensor must be used as an input for something else in the
+  graph. Otherwise, the autosummary op will not get executed, and the average
+  value will not get accumulated.
+- It is perfectly fine to include autosummaries with the same name in
+  several places throughout the graph, even if they are executed concurrently.
+- It is ok to also pass in a python scalar or numpy array. In this case, it
+  is added to the average immediately.
+"""
+
+from collections import OrderedDict
+import numpy as np
+import tensorflow as tf
+from tensorboard import summary as summary_lib
+from tensorboard.plugins.custom_scalar import layout_pb2
+
+from . import tfutil
+from .tfutil import TfExpression
+from .tfutil import TfExpressionEx
+
+_dtype = tf.float64
+_vars = OrderedDict()  # name => [var, ...]
+_immediate = OrderedDict()  # name => update_op, update_value
+_finalized = False
+_merge_op = None
+
+
+def _create_var(name: str, value_expr: TfExpression) -> TfExpression:
+    """Internal helper for creating autosummary accumulators."""
+    assert not _finalized
+    name_id = name.replace("/", "_")
+    v = tf.cast(value_expr, _dtype)
+
+    if v.shape.is_fully_defined():
+        size = np.prod(tfutil.shape_to_list(v.shape))
+        size_expr = tf.constant(size, dtype=_dtype)
+    else:
+        size = None
+        size_expr = tf.reduce_prod(tf.cast(tf.shape(v), _dtype))
+
+    if size == 1:
+        if v.shape.ndims != 0:
+            v = tf.reshape(v, [])
+        v = [size_expr, v, tf.square(v)]
+    else:
+        v = [size_expr, tf.reduce_sum(v), tf.reduce_sum(tf.square(v))]
+    v = tf.cond(tf.is_finite(v[1]), lambda: tf.stack(v), lambda: tf.zeros(3, dtype=_dtype))
+
+    with tfutil.absolute_name_scope("Autosummary/" + name_id), tf.control_dependencies(None):
+        var = tf.Variable(tf.zeros(3, dtype=_dtype), trainable=False)  # [sum(1), sum(x), sum(x**2)]
+    update_op = tf.cond(tf.is_variable_initialized(var), lambda: tf.assign_add(var, v), lambda: tf.assign(var, v))
+
+    if name in _vars:
+        _vars[name].append(var)
+    else:
+        _vars[name] = [var]
+    return update_op
+
+
+def autosummary(name: str, value: TfExpressionEx, passthru: TfExpressionEx = None) -> TfExpressionEx:
+    """Create a new autosummary.
+
+    Args:
+        name:     Name to use in TensorBoard
+        value:    TensorFlow expression or python value to track
+        passthru: Optionally return this TF node without modifications but tack an autosummary update side-effect to this node.
+
+    Example use of the passthru mechanism:
+
+    n = autosummary('l2loss', loss, passthru=n)
+
+    This is a shorthand for the following code:
+
+    with tf.control_dependencies([autosummary('l2loss', loss)]):
+        n = tf.identity(n)
+    """
+    tfutil.assert_tf_initialized()
+    name_id = name.replace("/", "_")
+
+    if tfutil.is_tf_expression(value):
+        with tf.name_scope("summary_" + name_id), tf.device(value.device):
+            update_op = _create_var(name, value)
+            with tf.control_dependencies([update_op]):
+                return tf.identity(value if passthru is None else passthru)
+
+    else:  # python scalar or numpy array
+        if name not in _immediate:
+            with tfutil.absolute_name_scope("Autosummary/" + name_id), tf.device(None), tf.control_dependencies(None):
+                update_value = tf.placeholder(_dtype)
+                update_op = _create_var(name, update_value)
+                _immediate[name] = update_op, update_value
+
+        update_op, update_value = _immediate[name]
+        tfutil.run(update_op, {update_value: value})
+        return value if passthru is None else passthru
+
+
+def finalize_autosummaries() -> None:
+    """Create the necessary ops to include autosummaries in TensorBoard report.
+    Note: This should be done only once per graph.
+    """
+    global _finalized
+    tfutil.assert_tf_initialized()
+
+    if _finalized:
+        return None
+
+    _finalized = True
+    tfutil.init_uninitialized_vars([var for vars_list in _vars.values() for var in vars_list])
+
+    # Create summary ops.
+    with tf.device(None), tf.control_dependencies(None):
+        for name, vars_list in _vars.items():
+            name_id = name.replace("/", "_")
+            with tfutil.absolute_name_scope("Autosummary/" + name_id):
+                moments = tf.add_n(vars_list)
+                moments /= moments[0]
+                with tf.control_dependencies([moments]):  # read before resetting
+                    reset_ops = [tf.assign(var, tf.zeros(3, dtype=_dtype)) for var in vars_list]
+                    with tf.name_scope(None), tf.control_dependencies(reset_ops):  # reset before reporting
+                        mean = moments[1]
+                        std = tf.sqrt(moments[2] - tf.square(moments[1]))
+                        tf.summary.scalar(name, mean)
+                        tf.summary.scalar("xCustomScalars/" + name + "/margin_lo", mean - std)
+                        tf.summary.scalar("xCustomScalars/" + name + "/margin_hi", mean + std)
+
+    # Group by category and chart name.
+    cat_dict = OrderedDict()
+    for series_name in sorted(_vars.keys()):
+        p = series_name.split("/")
+        cat = p[0] if len(p) >= 2 else ""
+        chart = "/".join(p[1:-1]) if len(p) >= 3 else p[-1]
+        if cat not in cat_dict:
+            cat_dict[cat] = OrderedDict()
+        if chart not in cat_dict[cat]:
+            cat_dict[cat][chart] = []
+        cat_dict[cat][chart].append(series_name)
+
+    # Setup custom_scalar layout.
+    categories = []
+    for cat_name, chart_dict in cat_dict.items():
+        charts = []
+        for chart_name, series_names in chart_dict.items():
+            series = []
+            for series_name in series_names:
+                series.append(layout_pb2.MarginChartContent.Series(
+                    value=series_name,
+                    lower="xCustomScalars/" + series_name + "/margin_lo",
+                    upper="xCustomScalars/" + series_name + "/margin_hi"))
+            margin = layout_pb2.MarginChartContent(series=series)
+            charts.append(layout_pb2.Chart(title=chart_name, margin=margin))
+        categories.append(layout_pb2.Category(title=cat_name, chart=charts))
+    layout = summary_lib.custom_scalar_pb(layout_pb2.Layout(category=categories))
+    return layout
+
+def save_summaries(file_writer, global_step=None):
+    """Call FileWriter.add_summary() with all summaries in the default graph,
+    automatically finalizing and merging them on the first call.
+    """
+    global _merge_op
+    tfutil.assert_tf_initialized()
+
+    if _merge_op is None:
+        layout = finalize_autosummaries()
+        if layout is not None:
+            file_writer.add_summary(layout)
+        with tf.device(None), tf.control_dependencies(None):
+            _merge_op = tf.summary.merge_all()
+
+    file_writer.add_summary(_merge_op.eval(), global_step)

dnnlib/tflib/network.py

@@ -0,0 +1,628 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial
+# 4.0 International License. To view a copy of this license, visit
+# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
+# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
+
+"""Helper for managing networks."""
+
+import types
+import inspect
+import re
+import uuid
+import sys
+import numpy as np
+import tensorflow as tf
+
+from collections import OrderedDict
+from typing import Any, List, Tuple, Union
+
+from . import tfutil
+from .. import util
+
+from .tfutil import TfExpression, TfExpressionEx
+
+_import_handlers = []  # Custom import handlers for dealing with legacy data in pickle import.
+_import_module_src = dict()  # Source code for temporary modules created during pickle import.
+
+
+def import_handler(handler_func):
+    """Function decorator for declaring custom import handlers."""
+    _import_handlers.append(handler_func)
+    return handler_func
+
+
+class Network:
+    """Generic network abstraction.
+
+    Acts as a convenience wrapper for a parameterized network construction
+    function, providing several utility methods and convenient access to
+    the inputs/outputs/weights.
+
+    Network objects can be safely pickled and unpickled for long-term
+    archival purposes. The pickling works reliably as long as the underlying
+    network construction function is defined in a standalone Python module
+    that has no side effects or application-specific imports.
+
+    Args:
+        name: Network name. Used to select TensorFlow name and variable scopes.
+        func_name: Fully qualified name of the underlying network construction function, or a top-level function object.
+        static_kwargs: Keyword arguments to be passed in to the network construction function.
+
+    Attributes:
+        name: User-specified name, defaults to build func name if None.
+        scope: Unique TensorFlow scope containing template graph and variables, derived from the user-specified name.
+        static_kwargs: Arguments passed to the user-supplied build func.
+        components: Container for sub-networks. Passed to the build func, and retained between calls.
+        num_inputs: Number of input tensors.
+        num_outputs: Number of output tensors.
+        input_shapes: Input tensor shapes (NC or NCHW), including minibatch dimension.
+        output_shapes: Output tensor shapes (NC or NCHW), including minibatch dimension.
+        input_shape: Short-hand for input_shapes[0].
+        output_shape: Short-hand for output_shapes[0].
+        input_templates: Input placeholders in the template graph.
+        output_templates: Output tensors in the template graph.
+        input_names: Name string for each input.
+        output_names: Name string for each output.
+        own_vars: Variables defined by this network (local_name => var), excluding sub-networks.
+        vars: All variables (local_name => var).
+        trainables: All trainable variables (local_name => var).
+        var_global_to_local: Mapping from variable global names to local names.
+    """
+
+    def __init__(self, name: str = None, func_name: Any = None, **static_kwargs):
+        tfutil.assert_tf_initialized()
+        assert isinstance(name, str) or name is None
+        assert func_name is not None
+        assert isinstance(func_name, str) or util.is_top_level_function(func_name)
+        assert util.is_pickleable(static_kwargs)
+
+        self._init_fields()
+        self.name = name
+        self.static_kwargs = util.EasyDict(static_kwargs)
+
+        # Locate the user-specified network build function.
+        if util.is_top_level_function(func_name):
+            func_name = util.get_top_level_function_name(func_name)
+        module, self._build_func_name = util.get_module_from_obj_name(func_name)
+        self._build_func = util.get_obj_from_module(module, self._build_func_name)
+        assert callable(self._build_func)
+
+        # Dig up source code for the module containing the build function.
+        self._build_module_src = _import_module_src.get(module, None)
+        if self._build_module_src is None:
+            self._build_module_src = inspect.getsource(module)
+
+        # Init TensorFlow graph.
+        self._init_graph()
+        self.reset_own_vars()
+
+    def _init_fields(self) -> None:
+        self.name = None
+        self.scope = None
+        self.static_kwargs = util.EasyDict()
+        self.components = util.EasyDict()
+        self.num_inputs = 0
+        self.num_outputs = 0
+        self.input_shapes = [[]]
+        self.output_shapes = [[]]
+        self.input_shape = []
+        self.output_shape = []
+        self.input_templates = []
+        self.output_templates = []
+        self.input_names = []
+        self.output_names = []
+        self.own_vars = OrderedDict()
+        self.vars = OrderedDict()
+        self.trainables = OrderedDict()
+        self.var_global_to_local = OrderedDict()
+
+        self._build_func = None  # User-supplied build function that constructs the network.
+        self._build_func_name = None  # Name of the build function.
+        self._build_module_src = None  # Full source code of the module containing the build function.
+        self._run_cache = dict()  # Cached graph data for Network.run().
+
+    def _init_graph(self) -> None:
+        # Collect inputs.
+        self.input_names = []
+
+        for param in inspect.signature(self._build_func).parameters.values():
+            if param.kind == param.POSITIONAL_OR_KEYWORD and param.default is param.empty:
+                self.input_names.append(param.name)
+
+        self.num_inputs = len(self.input_names)
+        assert self.num_inputs >= 1
+
+        # Choose name and scope.
+        if self.name is None:
+            self.name = self._build_func_name
+        assert re.match("^[A-Za-z0-9_.\\-]*$", self.name)
+        with tf.name_scope(None):
+            self.scope = tf.compat.v1.get_default_graph().unique_name(self.name, mark_as_used=True)
+
+        # Finalize build func kwargs.
+        build_kwargs = dict(self.static_kwargs)
+        build_kwargs["is_template_graph"] = True
+        build_kwargs["components"] = self.components
+
+        # Build template graph.
+        with tfutil.absolute_variable_scope(self.scope, reuse=tf.compat.v1.AUTO_REUSE), tfutil.absolute_name_scope(self.scope):  # ignore surrounding scopes
+            assert tf.compat.v1.get_variable_scope().name == self.scope
+            assert tf.compat.v1.get_default_graph().get_name_scope() == self.scope
+            with tf.control_dependencies(None):  # ignore surrounding control dependencies
+                self.input_templates = [tf.compat.v1.placeholder(tf.float32, name=name) for name in self.input_names]
+                out_expr = self._build_func(*self.input_templates, **build_kwargs)
+
+        # Collect outputs.
+        assert tfutil.is_tf_expression(out_expr) or isinstance(out_expr, tuple)
+        self.output_templates = [out_expr] if tfutil.is_tf_expression(out_expr) else list(out_expr)
+        self.num_outputs = len(self.output_templates)
+        assert self.num_outputs >= 1
+        assert all(tfutil.is_tf_expression(t) for t in self.output_templates)
+
+        # Perform sanity checks.
+        if any(t.shape.ndims is None for t in self.input_templates):
+            raise ValueError("Network input shapes not defined. Please call x.set_shape() for each input.")
+        if any(t.shape.ndims is None for t in self.output_templates):
+            raise ValueError("Network output shapes not defined. Please call x.set_shape() where applicable.")
+        if any(not isinstance(comp, Network) for comp in self.components.values()):
+            raise ValueError("Components of a Network must be Networks themselves.")
+        if len(self.components) != len(set(comp.name for comp in self.components.values())):
+            raise ValueError("Components of a Network must have unique names.")
+
+        # List inputs and outputs.
+        self.input_shapes = [tfutil.shape_to_list(t.shape) for t in self.input_templates]
+        self.output_shapes = [tfutil.shape_to_list(t.shape) for t in self.output_templates]
+        self.input_shape = self.input_shapes[0]
+        self.output_shape = self.output_shapes[0]
+        self.output_names = [t.name.split("/")[-1].split(":")[0] for t in self.output_templates]
+
+        # List variables.
+        self.own_vars = OrderedDict((var.name[len(self.scope) + 1:].split(":")[0], var) for var in tf.compat.v1.global_variables(self.scope + "/"))
+        self.vars = OrderedDict(self.own_vars)
+        self.vars.update((comp.name + "/" + name, var) for comp in self.components.values() for name, var in comp.vars.items())
+        self.trainables = OrderedDict((name, var) for name, var in self.vars.items() if var.trainable)
+        self.var_global_to_local = OrderedDict((var.name.split(":")[0], name) for name, var in self.vars.items())
+
+    def reset_own_vars(self) -> None:
+        """Re-initialize all variables of this network, excluding sub-networks."""
+        tfutil.run([var.initializer for var in self.own_vars.values()])
+
+    def reset_vars(self) -> None:
+        """Re-initialize all variables of this network, including sub-networks."""
+        tfutil.run([var.initializer for var in self.vars.values()])
+
+    def reset_trainables(self) -> None:
+        """Re-initialize all trainable variables of this network, including sub-networks."""
+        tfutil.run([var.initializer for var in self.trainables.values()])
+
+    def get_output_for(self, *in_expr: TfExpression, return_as_list: bool = False, **dynamic_kwargs) -> Union[TfExpression, List[TfExpression]]:
+        """Construct TensorFlow expression(s) for the output(s) of this network, given the input expression(s)."""
+        assert len(in_expr) == self.num_inputs
+        assert not all(expr is None for expr in in_expr)
+
+        # Finalize build func kwargs.
+        build_kwargs = dict(self.static_kwargs)
+        build_kwargs.update(dynamic_kwargs)
+        build_kwargs["is_template_graph"] = False
+        build_kwargs["components"] = self.components
+
+        # Build TensorFlow graph to evaluate the network.
+        with tfutil.absolute_variable_scope(self.scope, reuse=True), tf.name_scope(self.name):
+            assert tf.compat.v1.get_variable_scope().name == self.scope
+            valid_inputs = [expr for expr in in_expr if expr is not None]
+            final_inputs = []
+            for expr, name, shape in zip(in_expr, self.input_names, self.input_shapes):
+                if expr is not None:
+                    expr = tf.identity(expr, name=name)
+                else:
+                    expr = tf.zeros([tf.shape(valid_inputs[0])[0]] + shape[1:], name=name)
+                final_inputs.append(expr)
+            out_expr = self._build_func(*final_inputs, **build_kwargs)
+
+        # Propagate input shapes back to the user-specified expressions.
+        for expr, final in zip(in_expr, final_inputs):
+            if isinstance(expr, tf.Tensor):
+                expr.set_shape(final.shape)
+
+        # Express outputs in the desired format.
+        assert tfutil.is_tf_expression(out_expr) or isinstance(out_expr, tuple)
+        if return_as_list:
+            out_expr = [out_expr] if tfutil.is_tf_expression(out_expr) else list(out_expr)
+        return out_expr
+
+    def get_var_local_name(self, var_or_global_name: Union[TfExpression, str]) -> str:
+        """Get the local name of a given variable, without any surrounding name scopes."""
+        assert tfutil.is_tf_expression(var_or_global_name) or isinstance(var_or_global_name, str)
+        global_name = var_or_global_name if isinstance(var_or_global_name, str) else var_or_global_name.name
+        return self.var_global_to_local[global_name]
+
+    def find_var(self, var_or_local_name: Union[TfExpression, str]) -> TfExpression:
+        """Find variable by local or global name."""
+        assert tfutil.is_tf_expression(var_or_local_name) or isinstance(var_or_local_name, str)
+        return self.vars[var_or_local_name] if isinstance(var_or_local_name, str) else var_or_local_name
+
+    def get_var(self, var_or_local_name: Union[TfExpression, str]) -> np.ndarray:
+        """Get the value of a given variable as NumPy array.
+        Note: This method is very inefficient -- prefer to use tflib.run(list_of_vars) whenever possible."""
+        return self.find_var(var_or_local_name).eval()
+
+    def set_var(self, var_or_local_name: Union[TfExpression, str], new_value: Union[int, float, np.ndarray]) -> None:
+        """Set the value of a given variable based on the given NumPy array.
+        Note: This method is very inefficient -- prefer to use tflib.set_vars() whenever possible."""
+        tfutil.set_vars({self.find_var(var_or_local_name): new_value})
+
+    def __getstate__(self) -> dict:
+        """Pickle export."""
+        state = dict()
+        state["version"]            = 3
+        state["name"]               = self.name
+        state["static_kwargs"]      = dict(self.static_kwargs)
+        state["components"]         = dict(self.components)
+        state["build_module_src"]   = self._build_module_src
+        state["build_func_name"]    = self._build_func_name
+        state["variables"]          = list(zip(self.own_vars.keys(), tfutil.run(list(self.own_vars.values()))))
+        return state
+
+    def __setstate__(self, state: dict) -> None:
+        """Pickle import."""
+        # pylint: disable=attribute-defined-outside-init
+        tfutil.assert_tf_initialized()
+        self._init_fields()
+
+        # Execute custom import handlers.
+        for handler in _import_handlers:
+            state = handler(state)
+
+        # Set basic fields.
+        assert state["version"] in [2, 3]
+        self.name = state["name"]
+        self.static_kwargs = util.EasyDict(state["static_kwargs"])
+        self.components = util.EasyDict(state.get("components", {}))
+        self._build_module_src = state["build_module_src"]
+        self._build_func_name = state["build_func_name"]
+
+        # Create temporary module from the imported source code.
+        module_name = "_tflib_network_import_" + uuid.uuid4().hex
+        module = types.ModuleType(module_name)
+        sys.modules[module_name] = module
+        _import_module_src[module] = self._build_module_src
+        exec(self._build_module_src, module.__dict__) # pylint: disable=exec-used
+
+        # Locate network build function in the temporary module.
+        self._build_func = util.get_obj_from_module(module, self._build_func_name)
+        assert callable(self._build_func)
+
+        # Init TensorFlow graph.
+        self._init_graph()
+        self.reset_own_vars()
+        tfutil.set_vars({self.find_var(name): value for name, value in state["variables"]})
+
+    def clone(self, name: str = None, **new_static_kwargs) -> "Network":
+        """Create a clone of this network with its own copy of the variables."""
+        # pylint: disable=protected-access
+        net = object.__new__(Network)
+        net._init_fields()
+        net.name = name if name is not None else self.name
+        net.static_kwargs = util.EasyDict(self.static_kwargs)
+        net.static_kwargs.update(new_static_kwargs)
+        net._build_module_src = self._build_module_src
+        net._build_func_name = self._build_func_name
+        net._build_func = self._build_func
+        net._init_graph()
+        net.copy_vars_from(self)
+        return net
+
+    def copy_own_vars_from(self, src_net: "Network") -> None:
+        """Copy the values of all variables from the given network, excluding sub-networks."""
+        names = [name for name in self.own_vars.keys() if name in src_net.own_vars]
+        tfutil.set_vars(tfutil.run({self.vars[name]: src_net.vars[name] for name in names}))
+
+    def copy_vars_from(self, src_net: "Network") -> None:
+        """Copy the values of all variables from the given network, including sub-networks."""
+        names = [name for name in self.vars.keys() if name in src_net.vars]
+        tfutil.set_vars(tfutil.run({self.vars[name]: src_net.vars[name] for name in names}))
+
+    def copy_trainables_from(self, src_net: "Network") -> None:
+        """Copy the values of all trainable variables from the given network, including sub-networks."""
+        names = [name for name in self.trainables.keys() if name in src_net.trainables]
+        tfutil.set_vars(tfutil.run({self.vars[name]: src_net.vars[name] for name in names}))
+
+    def copy_compatible_trainables_from(self, src_net: "Network") -> None:
+        """Copy the compatible values of all trainable variables from the given network, including sub-networks"""
+        names = []
+        for name in self.trainables.keys():
+            if name not in src_net.trainables:
+                print("Not restoring (not present):     {}".format(name))
+            elif self.trainables[name].shape != src_net.trainables[name].shape:
+                print("Not restoring (different shape): {}".format(name))
+
+            if name in src_net.trainables and self.trainables[name].shape == src_net.trainables[name].shape:
+                names.append(name)
+
+        tfutil.set_vars(tfutil.run({self.vars[name]: src_net.vars[name] for name in names}))
+
+    def apply_swa(self, src_net, epoch):
+        """Perform stochastic weight averaging on the compatible values of all trainable variables from the given network, including sub-networks"""
+        names = []
+        for name in self.trainables.keys():
+            if name not in src_net.trainables:
+                print("Not restoring (not present):     {}".format(name))
+            elif self.trainables[name].shape != src_net.trainables[name].shape:
+                print("Not restoring (different shape): {}".format(name))
+
+            if name in src_net.trainables and self.trainables[name].shape == src_net.trainables[name].shape:
+                names.append(name)
+
+        scale_new_data = 1.0 / (epoch + 1)
+        scale_moving_average = (1.0 - scale_new_data)
+        tfutil.set_vars(tfutil.run({self.vars[name]: (src_net.vars[name] * scale_new_data + self.vars[name] * scale_moving_average) for name in names}))
+
+    def convert(self, new_func_name: str, new_name: str = None, **new_static_kwargs) -> "Network":
+        """Create new network with the given parameters, and copy all variables from this network."""
+        if new_name is None:
+            new_name = self.name
+        static_kwargs = dict(self.static_kwargs)
+        static_kwargs.update(new_static_kwargs)
+        net = Network(name=new_name, func_name=new_func_name, **static_kwargs)
+        net.copy_vars_from(self)
+        return net
+
+    def setup_as_moving_average_of(self, src_net: "Network", beta: TfExpressionEx = 0.99, beta_nontrainable: TfExpressionEx = 0.0) -> tf.Operation:
+        """Construct a TensorFlow op that updates the variables of this network
+        to be slightly closer to those of the given network."""
+        with tfutil.absolute_name_scope(self.scope + "/_MovingAvg"):
+            ops = []
+            for name, var in self.vars.items():
+                if name in src_net.vars:
+                    cur_beta = beta if name in self.trainables else beta_nontrainable
+                    new_value = tfutil.lerp(src_net.vars[name], var, cur_beta)
+                    ops.append(var.assign(new_value))
+            return tf.group(*ops)
+
+    def run(self,
+            *in_arrays: Tuple[Union[np.ndarray, None], ...],
+            input_transform: dict = None,
+            output_transform: dict = None,
+            return_as_list: bool = False,
+            print_progress: bool = False,
+            minibatch_size: int = None,
+            num_gpus: int = 1,
+            assume_frozen: bool = False,
+            custom_inputs=None,
+            **dynamic_kwargs) -> Union[np.ndarray, Tuple[np.ndarray, ...], List[np.ndarray]]:
+        """Run this network for the given NumPy array(s), and return the output(s) as NumPy array(s).
+
+        Args:
+            input_transform:    A dict specifying a custom transformation to be applied to the input tensor(s) before evaluating the network.
+                                The dict must contain a 'func' field that points to a top-level function. The function is called with the input
+                                TensorFlow expression(s) as positional arguments. Any remaining fields of the dict will be passed in as kwargs.
+            output_transform:   A dict specifying a custom transformation to be applied to the output tensor(s) after evaluating the network.
+                                The dict must contain a 'func' field that points to a top-level function. The function is called with the output
+                                TensorFlow expression(s) as positional arguments. Any remaining fields of the dict will be passed in as kwargs.
+            return_as_list:     True = return a list of NumPy arrays, False = return a single NumPy array, or a tuple if there are multiple outputs.
+            print_progress:     Print progress to the console? Useful for very large input arrays.
+            minibatch_size:     Maximum minibatch size to use, None = disable batching.
+            num_gpus:           Number of GPUs to use.
+            assume_frozen:      Improve multi-GPU performance by assuming that the trainable parameters will remain changed between calls.
+            dynamic_kwargs:     Additional keyword arguments to be passed into the network build function.
+            custom_inputs:      Allow to use another Tensor as input instead of default Placeholders
+        """
+        assert len(in_arrays) == self.num_inputs
+        assert not all(arr is None for arr in in_arrays)
+        assert input_transform is None or util.is_top_level_function(input_transform["func"])
+        assert output_transform is None or util.is_top_level_function(output_transform["func"])
+        output_transform, dynamic_kwargs = _handle_legacy_output_transforms(output_transform, dynamic_kwargs)
+        num_items = in_arrays[0].shape[0]
+        if minibatch_size is None:
+            minibatch_size = num_items
+
+        # Construct unique hash key from all arguments that affect the TensorFlow graph.
+        key = dict(input_transform=input_transform, output_transform=output_transform, num_gpus=num_gpus, assume_frozen=assume_frozen, dynamic_kwargs=dynamic_kwargs)
+        def unwind_key(obj):
+            if isinstance(obj, dict):
+                return [(key, unwind_key(value)) for key, value in sorted(obj.items())]
+            if callable(obj):
+                return util.get_top_level_function_name(obj)
+            return obj
+        key = repr(unwind_key(key))
+
+        # Build graph.
+        if key not in self._run_cache:
+            with tfutil.absolute_name_scope(self.scope + "/_Run"), tf.control_dependencies(None):
+                if custom_inputs is not None:
+                    with tf.device("/gpu:0"):
+                        in_expr = [input_builder(name) for input_builder, name in zip(custom_inputs, self.input_names)]
+                        in_split = list(zip(*[tf.split(x, num_gpus) for x in in_expr]))
+                else:
+                    with tf.device("/cpu:0"):
+                        in_expr = [tf.placeholder(tf.float32, name=name) for name in self.input_names]
+                        in_split = list(zip(*[tf.split(x, num_gpus) for x in in_expr]))
+
+                out_split = []
+                for gpu in range(num_gpus):
+                    with tf.device("/gpu:%d" % gpu):
+                        net_gpu = self.clone() if assume_frozen else self
+                        in_gpu = in_split[gpu]
+
+                        if input_transform is not None:
+                            in_kwargs = dict(input_transform)
+                            in_gpu = in_kwargs.pop("func")(*in_gpu, **in_kwargs)
+                            in_gpu = [in_gpu] if tfutil.is_tf_expression(in_gpu) else list(in_gpu)
+
+                        assert len(in_gpu) == self.num_inputs
+                        out_gpu = net_gpu.get_output_for(*in_gpu, return_as_list=True, **dynamic_kwargs)
+
+                        if output_transform is not None:
+                            out_kwargs = dict(output_transform)
+                            out_gpu = out_kwargs.pop("func")(*out_gpu, **out_kwargs)
+                            out_gpu = [out_gpu] if tfutil.is_tf_expression(out_gpu) else list(out_gpu)
+
+                        assert len(out_gpu) == self.num_outputs
+                        out_split.append(out_gpu)
+
+                with tf.device("/cpu:0"):
+                    out_expr = [tf.concat(outputs, axis=0) for outputs in zip(*out_split)]
+                    self._run_cache[key] = in_expr, out_expr
+
+        # Run minibatches.
+        in_expr, out_expr = self._run_cache[key]
+        out_arrays = [np.empty([num_items] + tfutil.shape_to_list(expr.shape)[1:], expr.dtype.name) for expr in out_expr]
+
+        for mb_begin in range(0, num_items, minibatch_size):
+            if print_progress:
+                print("\r%d / %d" % (mb_begin, num_items), end="")
+
+            mb_end = min(mb_begin + minibatch_size, num_items)
+            mb_num = mb_end - mb_begin
+            mb_in = [src[mb_begin : mb_end] if src is not None else np.zeros([mb_num] + shape[1:]) for src, shape in zip(in_arrays, self.input_shapes)]
+            mb_out = tf.compat.v1.get_default_session().run(out_expr, dict(zip(in_expr, mb_in)))
+
+            for dst, src in zip(out_arrays, mb_out):
+                dst[mb_begin: mb_end] = src
+
+        # Done.
+        if print_progress:
+            print("\r%d / %d" % (num_items, num_items))
+
+        if not return_as_list:
+            out_arrays = out_arrays[0] if len(out_arrays) == 1 else tuple(out_arrays)
+        return out_arrays
+
+    def list_ops(self) -> List[TfExpression]:
+        include_prefix = self.scope + "/"
+        exclude_prefix = include_prefix + "_"
+        ops = tf.get_default_graph().get_operations()
+        ops = [op for op in ops if op.name.startswith(include_prefix)]
+        ops = [op for op in ops if not op.name.startswith(exclude_prefix)]
+        return ops
+
+    def list_layers(self) -> List[Tuple[str, TfExpression, List[TfExpression]]]:
+        """Returns a list of (layer_name, output_expr, trainable_vars) tuples corresponding to
+        individual layers of the network. Mainly intended to be used for reporting."""
+        layers = []
+
+        def recurse(scope, parent_ops, parent_vars, level):
+            # Ignore specific patterns.
+            if any(p in scope for p in ["/Shape", "/strided_slice", "/Cast", "/concat", "/Assign"]):
+                return
+
+            # Filter ops and vars by scope.
+            global_prefix = scope + "/"
+            local_prefix = global_prefix[len(self.scope) + 1:]
+            cur_ops = [op for op in parent_ops if op.name.startswith(global_prefix) or op.name == global_prefix[:-1]]
+            cur_vars = [(name, var) for name, var in parent_vars if name.startswith(local_prefix) or name == local_prefix[:-1]]
+            if not cur_ops and not cur_vars:
+                return
+
+            # Filter out all ops related to variables.
+            for var in [op for op in cur_ops if op.type.startswith("Variable")]:
+                var_prefix = var.name + "/"
+                cur_ops = [op for op in cur_ops if not op.name.startswith(var_prefix)]
+
+            # Scope does not contain ops as immediate children => recurse deeper.
+            contains_direct_ops = any("/" not in op.name[len(global_prefix):] and op.type != "Identity" for op in cur_ops)
+            if (level == 0 or not contains_direct_ops) and (len(cur_ops) + len(cur_vars)) > 1:
+                visited = set()
+                for rel_name in [op.name[len(global_prefix):] for op in cur_ops] + [name[len(local_prefix):] for name, _var in cur_vars]:
+                    token = rel_name.split("/")[0]
+                    if token not in visited:
+                        recurse(global_prefix + token, cur_ops, cur_vars, level + 1)
+                        visited.add(token)
+                return
+
+            # Report layer.
+            layer_name = scope[len(self.scope) + 1:]
+            layer_output = cur_ops[-1].outputs[0] if cur_ops else cur_vars[-1][1]
+            layer_trainables = [var for _name, var in cur_vars if var.trainable]
+            layers.append((layer_name, layer_output, layer_trainables))
+
+        recurse(self.scope, self.list_ops(), list(self.vars.items()), 0)
+        return layers
+
+    def print_layers(self, title: str = None, hide_layers_with_no_params: bool = False) -> None:
+        """Print a summary table of the network structure."""
+        rows = [[title if title is not None else self.name, "Params", "OutputShape", "WeightShape"]]
+        rows += [["---"] * 4]
+        total_params = 0
+
+        for layer_name, layer_output, layer_trainables in self.list_layers():
+            num_params = sum(np.prod(tfutil.shape_to_list(var.shape)) for var in layer_trainables)
+            weights = [var for var in layer_trainables if var.name.endswith("/weight:0") or var.name.endswith("/weight_1:0")]
+            weights.sort(key=lambda x: len(x.name))
+            if len(weights) == 0 and len(layer_trainables) == 1:
+                weights = layer_trainables
+            total_params += num_params
+
+            if not hide_layers_with_no_params or num_params != 0:
+                num_params_str = str(num_params) if num_params > 0 else "-"
+                output_shape_str = str(layer_output.shape)
+                weight_shape_str = str(weights[0].shape) if len(weights) >= 1 else "-"
+                rows += [[layer_name, num_params_str, output_shape_str, weight_shape_str]]
+
+        rows += [["---"] * 4]
+        rows += [["Total", str(total_params), "", ""]]
+
+        widths = [max(len(cell) for cell in column) for column in zip(*rows)]
+        print()
+        for row in rows:
+            print("  ".join(cell + " " * (width - len(cell)) for cell, width in zip(row, widths)))
+        print()
+
+    def setup_weight_histograms(self, title: str = None) -> None:
+        """Construct summary ops to include histograms of all trainable parameters in TensorBoard."""
+        if title is None:
+            title = self.name
+
+        with tf.name_scope(None), tf.device(None), tf.control_dependencies(None):
+            for local_name, var in self.trainables.items():
+                if "/" in local_name:
+                    p = local_name.split("/")
+                    name = title + "_" + p[-1] + "/" + "_".join(p[:-1])
+                else:
+                    name = title + "_toplevel/" + local_name
+
+                tf.summary.histogram(name, var)
+
+#----------------------------------------------------------------------------
+# Backwards-compatible emulation of legacy output transformation in Network.run().
+
+_print_legacy_warning = True
+
+def _handle_legacy_output_transforms(output_transform, dynamic_kwargs):
+    global _print_legacy_warning
+    legacy_kwargs = ["out_mul", "out_add", "out_shrink", "out_dtype"]
+    if not any(kwarg in dynamic_kwargs for kwarg in legacy_kwargs):
+        return output_transform, dynamic_kwargs
+
+    if _print_legacy_warning:
+        _print_legacy_warning = False
+        print()
+        print("WARNING: Old-style output transformations in Network.run() are deprecated.")
+        print("Consider using 'output_transform=dict(func=tflib.convert_images_to_uint8)'")
+        print("instead of 'out_mul=127.5, out_add=127.5, out_dtype=np.uint8'.")
+        print()
+    assert output_transform is None
+
+    new_kwargs = dict(dynamic_kwargs)
+    new_transform = {kwarg: new_kwargs.pop(kwarg) for kwarg in legacy_kwargs if kwarg in dynamic_kwargs}
+    new_transform["func"] = _legacy_output_transform_func
+    return new_transform, new_kwargs
+
+def _legacy_output_transform_func(*expr, out_mul=1.0, out_add=0.0, out_shrink=1, out_dtype=None):
+    if out_mul != 1.0:
+        expr = [x * out_mul for x in expr]
+
+    if out_add != 0.0:
+        expr = [x + out_add for x in expr]
+
+    if out_shrink > 1:
+        ksize = [1, 1, out_shrink, out_shrink]
+        expr = [tf.nn.avg_pool(x, ksize=ksize, strides=ksize, padding="VALID", data_format="NCHW") for x in expr]
+
+    if out_dtype is not None:
+        if tf.as_dtype(out_dtype).is_integer:
+            expr = [tf.round(x) for x in expr]
+        expr = [tf.saturate_cast(x, out_dtype) for x in expr]
+    return expr

dnnlib/tflib/optimizer.py

@@ -0,0 +1,214 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial
+# 4.0 International License. To view a copy of this license, visit
+# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
+# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
+
+"""Helper wrapper for a Tensorflow optimizer."""
+
+import numpy as np
+import tensorflow as tf
+
+from collections import OrderedDict
+from typing import List, Union
+
+from . import autosummary
+from . import tfutil
+from .. import util
+
+from .tfutil import TfExpression, TfExpressionEx
+
+try:
+    # TensorFlow 1.13
+    from tensorflow.python.ops import nccl_ops
+except:
+    # Older TensorFlow versions
+    import tensorflow.contrib.nccl as nccl_ops
+
+class Optimizer:
+    """A Wrapper for tf.train.Optimizer.
+
+    Automatically takes care of:
+    - Gradient averaging for multi-GPU training.
+    - Dynamic loss scaling and typecasts for FP16 training.
+    - Ignoring corrupted gradients that contain NaNs/Infs.
+    - Reporting statistics.
+    - Well-chosen default settings.
+    """
+
+    def __init__(self,
+                 name: str = "Train",
+                 tf_optimizer: str = "tf.train.AdamOptimizer",
+                 learning_rate: TfExpressionEx = 0.001,
+                 use_loss_scaling: bool = False,
+                 loss_scaling_init: float = 64.0,
+                 loss_scaling_inc: float = 0.0005,
+                 loss_scaling_dec: float = 1.0,
+                 **kwargs):
+
+        # Init fields.
+        self.name = name
+        self.learning_rate = tf.convert_to_tensor(learning_rate)
+        self.id = self.name.replace("/", ".")
+        self.scope = tf.get_default_graph().unique_name(self.id)
+        self.optimizer_class = util.get_obj_by_name(tf_optimizer)
+        self.optimizer_kwargs = dict(kwargs)
+        self.use_loss_scaling = use_loss_scaling
+        self.loss_scaling_init = loss_scaling_init
+        self.loss_scaling_inc = loss_scaling_inc
+        self.loss_scaling_dec = loss_scaling_dec
+        self._grad_shapes = None  # [shape, ...]
+        self._dev_opt = OrderedDict()  # device => optimizer
+        self._dev_grads = OrderedDict()  # device => [[(grad, var), ...], ...]
+        self._dev_ls_var = OrderedDict()  # device => variable (log2 of loss scaling factor)
+        self._updates_applied = False
+
+    def register_gradients(self, loss: TfExpression, trainable_vars: Union[List, dict]) -> None:
+        """Register the gradients of the given loss function with respect to the given variables.
+        Intended to be called once per GPU."""
+        assert not self._updates_applied
+
+        # Validate arguments.
+        if isinstance(trainable_vars, dict):
+            trainable_vars = list(trainable_vars.values())  # allow passing in Network.trainables as vars
+
+        assert isinstance(trainable_vars, list) and len(trainable_vars) >= 1
+        assert all(tfutil.is_tf_expression(expr) for expr in trainable_vars + [loss])
+
+        if self._grad_shapes is None:
+            self._grad_shapes = [tfutil.shape_to_list(var.shape) for var in trainable_vars]
+
+        assert len(trainable_vars) == len(self._grad_shapes)
+        assert all(tfutil.shape_to_list(var.shape) == var_shape for var, var_shape in zip(trainable_vars, self._grad_shapes))
+
+        dev = loss.device
+
+        assert all(var.device == dev for var in trainable_vars)
+
+        # Register device and compute gradients.
+        with tf.name_scope(self.id + "_grad"), tf.device(dev):
+            if dev not in self._dev_opt:
+                opt_name = self.scope.replace("/", "_") + "_opt%d" % len(self._dev_opt)
+                assert callable(self.optimizer_class)
+                self._dev_opt[dev] = self.optimizer_class(name=opt_name, learning_rate=self.learning_rate, **self.optimizer_kwargs)
+                self._dev_grads[dev] = []
+
+            loss = self.apply_loss_scaling(tf.cast(loss, tf.float32))
+            grads = self._dev_opt[dev].compute_gradients(loss, trainable_vars, gate_gradients=tf.train.Optimizer.GATE_NONE)  # disable gating to reduce memory usage
+            grads = [(g, v) if g is not None else (tf.zeros_like(v), v) for g, v in grads]  # replace disconnected gradients with zeros
+            self._dev_grads[dev].append(grads)
+
+    def apply_updates(self) -> tf.Operation:
+        """Construct training op to update the registered variables based on their gradients."""
+        tfutil.assert_tf_initialized()
+        assert not self._updates_applied
+        self._updates_applied = True
+        devices = list(self._dev_grads.keys())
+        total_grads = sum(len(grads) for grads in self._dev_grads.values())
+        assert len(devices) >= 1 and total_grads >= 1
+        ops = []
+
+        with tfutil.absolute_name_scope(self.scope):
+            # Cast gradients to FP32 and calculate partial sum within each device.
+            dev_grads = OrderedDict()  # device => [(grad, var), ...]
+
+            for dev_idx, dev in enumerate(devices):
+                with tf.name_scope("ProcessGrads%d" % dev_idx), tf.device(dev):
+                    sums = []
+
+                    for gv in zip(*self._dev_grads[dev]):
+                        assert all(v is gv[0][1] for g, v in gv)
+                        g = [tf.cast(g, tf.float32) for g, v in gv]
+                        g = g[0] if len(g) == 1 else tf.add_n(g)
+                        sums.append((g, gv[0][1]))
+
+                    dev_grads[dev] = sums
+
+            # Sum gradients across devices.
+            if len(devices) > 1:
+                with tf.name_scope("SumAcrossGPUs"), tf.device(None):
+                    for var_idx, grad_shape in enumerate(self._grad_shapes):
+                        g = [dev_grads[dev][var_idx][0] for dev in devices]
+
+                        if np.prod(grad_shape):  # nccl does not support zero-sized tensors
+                            g = nccl_ops.all_sum(g)
+
+                        for dev, gg in zip(devices, g):
+                            dev_grads[dev][var_idx] = (gg, dev_grads[dev][var_idx][1])
+
+            # Apply updates separately on each device.
+            for dev_idx, (dev, grads) in enumerate(dev_grads.items()):
+                with tf.name_scope("ApplyGrads%d" % dev_idx), tf.device(dev):
+                    # Scale gradients as needed.
+                    if self.use_loss_scaling or total_grads > 1:
+                        with tf.name_scope("Scale"):
+                            coef = tf.constant(np.float32(1.0 / total_grads), name="coef")
+                            coef = self.undo_loss_scaling(coef)
+                            grads = [(g * coef, v) for g, v in grads]
+
+                    # Check for overflows.
+                    with tf.name_scope("CheckOverflow"):
+                        grad_ok = tf.reduce_all(tf.stack([tf.reduce_all(tf.is_finite(g)) for g, v in grads]))
+
+                    # Update weights and adjust loss scaling.
+                    with tf.name_scope("UpdateWeights"):
+                        # pylint: disable=cell-var-from-loop
+                        opt = self._dev_opt[dev]
+                        ls_var = self.get_loss_scaling_var(dev)
+
+                        if not self.use_loss_scaling:
+                            ops.append(tf.cond(grad_ok, lambda: opt.apply_gradients(grads), tf.no_op))
+                        else:
+                            ops.append(tf.cond(grad_ok,
+                                               lambda: tf.group(tf.assign_add(ls_var, self.loss_scaling_inc), opt.apply_gradients(grads)),
+                                               lambda: tf.group(tf.assign_sub(ls_var, self.loss_scaling_dec))))
+
+                    # Report statistics on the last device.
+                    if dev == devices[-1]:
+                        with tf.name_scope("Statistics"):
+                            ops.append(autosummary.autosummary(self.id + "/learning_rate", self.learning_rate))
+                            ops.append(autosummary.autosummary(self.id + "/overflow_frequency", tf.where(grad_ok, 0, 1)))
+
+                            if self.use_loss_scaling:
+                                ops.append(autosummary.autosummary(self.id + "/loss_scaling_log2", ls_var))
+
+            # Initialize variables and group everything into a single op.
+            self.reset_optimizer_state()
+            tfutil.init_uninitialized_vars(list(self._dev_ls_var.values()))
+
+            return tf.group(*ops, name="TrainingOp")
+
+    def reset_optimizer_state(self) -> None:
+        """Reset internal state of the underlying optimizer."""
+        tfutil.assert_tf_initialized()
+        tfutil.run([var.initializer for opt in self._dev_opt.values() for var in opt.variables()])
+
+    def get_loss_scaling_var(self, device: str) -> Union[tf.Variable, None]:
+        """Get or create variable representing log2 of the current dynamic loss scaling factor."""
+        if not self.use_loss_scaling:
+            return None
+
+        if device not in self._dev_ls_var:
+            with tfutil.absolute_name_scope(self.scope + "/LossScalingVars"), tf.control_dependencies(None):
+                self._dev_ls_var[device] = tf.Variable(np.float32(self.loss_scaling_init), name="loss_scaling_var")
+
+        return self._dev_ls_var[device]
+
+    def apply_loss_scaling(self, value: TfExpression) -> TfExpression:
+        """Apply dynamic loss scaling for the given expression."""
+        assert tfutil.is_tf_expression(value)
+
+        if not self.use_loss_scaling:
+            return value
+
+        return value * tfutil.exp2(self.get_loss_scaling_var(value.device))
+
+    def undo_loss_scaling(self, value: TfExpression) -> TfExpression:
+        """Undo the effect of dynamic loss scaling for the given expression."""
+        assert tfutil.is_tf_expression(value)
+
+        if not self.use_loss_scaling:
+            return value
+
+        return value * tfutil.exp2(-self.get_loss_scaling_var(value.device)) # pylint: disable=invalid-unary-operand-type

dnnlib/tflib/tfutil.py

@@ -0,0 +1,242 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial
+# 4.0 International License. To view a copy of this license, visit
+# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
+# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
+
+"""Miscellaneous helper utils for Tensorflow."""
+
+import os
+import numpy as np
+import tensorflow as tf
+
+from typing import Any, Iterable, List, Union
+
+TfExpression = Union[tf.Tensor, tf.Variable, tf.Operation]
+"""A type that represents a valid Tensorflow expression."""
+
+TfExpressionEx = Union[TfExpression, int, float, np.ndarray]
+"""A type that can be converted to a valid Tensorflow expression."""
+
+
+def run(*args, **kwargs) -> Any:
+    """Run the specified ops in the default session."""
+    assert_tf_initialized()
+    return tf.compat.v1.get_default_session().run(*args, **kwargs)
+
+
+def is_tf_expression(x: Any) -> bool:
+    """Check whether the input is a valid Tensorflow expression, i.e., Tensorflow Tensor, Variable, or Operation."""
+    return isinstance(x, (tf.Tensor, tf.Variable, tf.Operation))
+
+
+def shape_to_list(shape: Iterable[tf.compat.v1.Dimension]) -> List[Union[int, None]]:
+    """Convert a Tensorflow shape to a list of ints."""
+    return [dim.value for dim in shape]
+
+
+def flatten(x: TfExpressionEx) -> TfExpression:
+    """Shortcut function for flattening a tensor."""
+    with tf.name_scope("Flatten"):
+        return tf.reshape(x, [-1])
+
+
+def log2(x: TfExpressionEx) -> TfExpression:
+    """Logarithm in base 2."""
+    with tf.name_scope("Log2"):
+        return tf.log(x) * np.float32(1.0 / np.log(2.0))
+
+
+def exp2(x: TfExpressionEx) -> TfExpression:
+    """Exponent in base 2."""
+    with tf.name_scope("Exp2"):
+        return tf.exp(x * np.float32(np.log(2.0)))
+
+
+def lerp(a: TfExpressionEx, b: TfExpressionEx, t: TfExpressionEx) -> TfExpressionEx:
+    """Linear interpolation."""
+    with tf.name_scope("Lerp"):
+        return a + (b - a) * t
+
+
+def lerp_clip(a: TfExpressionEx, b: TfExpressionEx, t: TfExpressionEx) -> TfExpression:
+    """Linear interpolation with clip."""
+    with tf.name_scope("LerpClip"):
+        return a + (b - a) * tf.clip_by_value(t, 0.0, 1.0)
+
+
+def absolute_name_scope(scope: str) -> tf.name_scope:
+    """Forcefully enter the specified name scope, ignoring any surrounding scopes."""
+    return tf.name_scope(scope + "/")
+
+
+def absolute_variable_scope(scope: str, **kwargs) -> tf.compat.v1.variable_scope:
+    """Forcefully enter the specified variable scope, ignoring any surrounding scopes."""
+    return tf.compat.v1.variable_scope(tf.compat.v1.VariableScope(name=scope, **kwargs), auxiliary_name_scope=False)
+
+
+def _sanitize_tf_config(config_dict: dict = None) -> dict:
+    # Defaults.
+    cfg = dict()
+    cfg["rnd.np_random_seed"]               = None      # Random seed for NumPy. None = keep as is.
+    cfg["rnd.tf_random_seed"]               = "auto"    # Random seed for TensorFlow. 'auto' = derive from NumPy random state. None = keep as is.
+    cfg["env.TF_CPP_MIN_LOG_LEVEL"]         = "1"       # 0 = Print all available debug info from TensorFlow. 1 = Print warnings and errors, but disable debug info.
+    cfg["graph_options.place_pruned_graph"] = True      # False = Check that all ops are available on the designated device. True = Skip the check for ops that are not used.
+    cfg["gpu_options.allow_growth"]         = True      # False = Allocate all GPU memory at the beginning. True = Allocate only as much GPU memory as needed.
+
+    # User overrides.
+    if config_dict is not None:
+        cfg.update(config_dict)
+    return cfg
+
+
+def init_tf(config_dict: dict = None) -> None:
+    """Initialize TensorFlow session using good default settings."""
+    # Skip if already initialized.
+    if tf.compat.v1.get_default_session() is not None:
+        return
+
+    # Setup config dict and random seeds.
+    cfg = _sanitize_tf_config(config_dict)
+    np_random_seed = cfg["rnd.np_random_seed"]
+    if np_random_seed is not None:
+        np.random.seed(np_random_seed)
+    tf_random_seed = cfg["rnd.tf_random_seed"]
+    if tf_random_seed == "auto":
+        tf_random_seed = np.random.randint(1 << 31)
+    if tf_random_seed is not None:
+        tf.compat.v1.set_random_seed(tf_random_seed)
+
+    # Setup environment variables.
+    for key, value in list(cfg.items()):
+        fields = key.split(".")
+        if fields[0] == "env":
+            assert len(fields) == 2
+            os.environ[fields[1]] = str(value)
+
+    # Create default TensorFlow session.
+    create_session(cfg, force_as_default=True)
+
+
+def assert_tf_initialized():
+    """Check that TensorFlow session has been initialized."""
+    if tf.compat.v1.get_default_session() is None:
+        raise RuntimeError("No default TensorFlow session found. Please call dnnlib.tflib.init_tf().")
+
+
+def create_session(config_dict: dict = None, force_as_default: bool = False) -> tf.compat.v1.Session:
+    """Create tf.Session based on config dict."""
+    # Setup TensorFlow config proto.
+    cfg = _sanitize_tf_config(config_dict)
+    config_proto = tf.compat.v1.ConfigProto()
+    for key, value in cfg.items():
+        fields = key.split(".")
+        if fields[0] not in ["rnd", "env"]:
+            obj = config_proto
+            for field in fields[:-1]:
+                obj = getattr(obj, field)
+            setattr(obj, fields[-1], value)
+
+    # Create session.
+    session = tf.compat.v1.Session(config=config_proto)
+    if force_as_default:
+        # pylint: disable=protected-access
+        session._default_session = session.as_default()
+        session._default_session.enforce_nesting = False
+        session._default_session.__enter__() # pylint: disable=no-member
+
+    return session
+
+
+def init_uninitialized_vars(target_vars: List[tf.Variable] = None) -> None:
+    """Initialize all tf.Variables that have not already been initialized.
+
+    Equivalent to the following, but more efficient and does not bloat the tf graph:
+    tf.variables_initializer(tf.report_uninitialized_variables()).run()
+    """
+    assert_tf_initialized()
+    if target_vars is None:
+        target_vars = tf.global_variables()
+
+    test_vars = []
+    test_ops = []
+
+    with tf.control_dependencies(None):  # ignore surrounding control_dependencies
+        for var in target_vars:
+            assert is_tf_expression(var)
+
+            try:
+                tf.get_default_graph().get_tensor_by_name(var.name.replace(":0", "/IsVariableInitialized:0"))
+            except KeyError:
+                # Op does not exist => variable may be uninitialized.
+                test_vars.append(var)
+
+                with absolute_name_scope(var.name.split(":")[0]):
+                    test_ops.append(tf.is_variable_initialized(var))
+
+    init_vars = [var for var, inited in zip(test_vars, run(test_ops)) if not inited]
+    run([var.initializer for var in init_vars])
+
+
+def set_vars(var_to_value_dict: dict) -> None:
+    """Set the values of given tf.Variables.
+
+    Equivalent to the following, but more efficient and does not bloat the tf graph:
+    tflib.run([tf.assign(var, value) for var, value in var_to_value_dict.items()]
+    """
+    assert_tf_initialized()
+    ops = []
+    feed_dict = {}
+
+    for var, value in var_to_value_dict.items():
+        assert is_tf_expression(var)
+
+        try:
+            setter = tf.compat.v1.get_default_graph().get_tensor_by_name(var.name.replace(":0", "/setter:0"))  # look for existing op
+        except KeyError:
+            with absolute_name_scope(var.name.split(":")[0]):
+                with tf.control_dependencies(None):  # ignore surrounding control_dependencies
+                    setter = tf.compat.v1.assign(var, tf.compat.v1.placeholder(var.dtype, var.shape, "new_value"), name="setter")  # create new setter
+
+        ops.append(setter)
+        feed_dict[setter.op.inputs[1]] = value
+
+    run(ops, feed_dict)
+
+
+def create_var_with_large_initial_value(initial_value: np.ndarray, *args, **kwargs):
+    """Create tf.Variable with large initial value without bloating the tf graph."""
+    assert_tf_initialized()
+    assert isinstance(initial_value, np.ndarray)
+    zeros = tf.zeros(initial_value.shape, initial_value.dtype)
+    var = tf.Variable(zeros, *args, **kwargs)
+    set_vars({var: initial_value})
+    return var
+
+
+def convert_images_from_uint8(images, drange=[-1,1], nhwc_to_nchw=False):
+    """Convert a minibatch of images from uint8 to float32 with configurable dynamic range.
+    Can be used as an input transformation for Network.run().
+    """
+    images = tf.cast(images, tf.float32)
+    if nhwc_to_nchw:
+        images = tf.transpose(images, [0, 3, 1, 2])
+    return (images - drange[0]) * ((drange[1] - drange[0]) / 255)
+
+
+def convert_images_to_uint8(images, drange=[-1,1], nchw_to_nhwc=False, shrink=1, uint8_cast=True):
+    """Convert a minibatch of images from float32 to uint8 with configurable dynamic range.
+    Can be used as an output transformation for Network.run().
+    """
+    images = tf.cast(images, tf.float32)
+    if shrink > 1:
+        ksize = [1, 1, shrink, shrink]
+        images = tf.nn.avg_pool(images, ksize=ksize, strides=ksize, padding="VALID", data_format="NCHW")
+    if nchw_to_nhwc:
+        images = tf.transpose(images, [0, 2, 3, 1])
+    scale = 255 / (drange[1] - drange[0])
+    images = images * scale + (0.5 - drange[0] * scale)
+    if uint8_cast:
+        images = tf.saturate_cast(images, tf.uint8)
+    return images

dnnlib/util.py

@@ -0,0 +1,408 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial
+# 4.0 International License. To view a copy of this license, visit
+# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
+# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
+
+"""Miscellaneous utility classes and functions."""
+
+import ctypes
+import fnmatch
+import importlib
+import inspect
+import numpy as np
+import os
+import shutil
+import sys
+import types
+import io
+import pickle
+import re
+import requests
+import html
+import hashlib
+import glob
+import uuid
+
+from distutils.util import strtobool
+from typing import Any, List, Tuple, Union
+
+
+# Util classes
+# ------------------------------------------------------------------------------------------
+
+
+class EasyDict(dict):
+    """Convenience class that behaves like a dict but allows access with the attribute syntax."""
+
+    def __getattr__(self, name: str) -> Any:
+        try:
+            return self[name]
+        except KeyError:
+            raise AttributeError(name)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        self[name] = value
+
+    def __delattr__(self, name: str) -> None:
+        del self[name]
+
+
+class Logger(object):
+    """Redirect stderr to stdout, optionally print stdout to a file, and optionally force flushing on both stdout and the file."""
+
+    def __init__(self, file_name: str = None, file_mode: str = "w", should_flush: bool = True):
+        self.file = None
+
+        if file_name is not None:
+            self.file = open(file_name, file_mode)
+
+        self.should_flush = should_flush
+        self.stdout = sys.stdout
+        self.stderr = sys.stderr
+
+        sys.stdout = self
+        sys.stderr = self
+
+    def __enter__(self) -> "Logger":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.close()
+
+    def write(self, text: str) -> None:
+        """Write text to stdout (and a file) and optionally flush."""
+        if len(text) == 0: # workaround for a bug in VSCode debugger: sys.stdout.write(''); sys.stdout.flush() => crash
+            return
+
+        if self.file is not None:
+            self.file.write(text)
+
+        self.stdout.write(text)
+
+        if self.should_flush:
+            self.flush()
+
+    def flush(self) -> None:
+        """Flush written text to both stdout and a file, if open."""
+        if self.file is not None:
+            self.file.flush()
+
+        self.stdout.flush()
+
+    def close(self) -> None:
+        """Flush, close possible files, and remove stdout/stderr mirroring."""
+        self.flush()
+
+        # if using multiple loggers, prevent closing in wrong order
+        if sys.stdout is self:
+            sys.stdout = self.stdout
+        if sys.stderr is self:
+            sys.stderr = self.stderr
+
+        if self.file is not None:
+            self.file.close()
+
+
+# Small util functions
+# ------------------------------------------------------------------------------------------
+
+
+def format_time(seconds: Union[int, float]) -> str:
+    """Convert the seconds to human readable string with days, hours, minutes and seconds."""
+    s = int(np.rint(seconds))
+
+    if s < 60:
+        return "{0}s".format(s)
+    elif s < 60 * 60:
+        return "{0}m {1:02}s".format(s // 60, s % 60)
+    elif s < 24 * 60 * 60:
+        return "{0}h {1:02}m {2:02}s".format(s // (60 * 60), (s // 60) % 60, s % 60)
+    else:
+        return "{0}d {1:02}h {2:02}m".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24, (s // 60) % 60)
+
+
+def ask_yes_no(question: str) -> bool:
+    """Ask the user the question until the user inputs a valid answer."""
+    while True:
+        try:
+            print("{0} [y/n]".format(question))
+            return strtobool(input().lower())
+        except ValueError:
+            pass
+
+
+def tuple_product(t: Tuple) -> Any:
+    """Calculate the product of the tuple elements."""
+    result = 1
+
+    for v in t:
+        result *= v
+
+    return result
+
+
+_str_to_ctype = {
+    "uint8": ctypes.c_ubyte,
+    "uint16": ctypes.c_uint16,
+    "uint32": ctypes.c_uint32,
+    "uint64": ctypes.c_uint64,
+    "int8": ctypes.c_byte,
+    "int16": ctypes.c_int16,
+    "int32": ctypes.c_int32,
+    "int64": ctypes.c_int64,
+    "float32": ctypes.c_float,
+    "float64": ctypes.c_double
+}
+
+
+def get_dtype_and_ctype(type_obj: Any) -> Tuple[np.dtype, Any]:
+    """Given a type name string (or an object having a __name__ attribute), return matching Numpy and ctypes types that have the same size in bytes."""
+    type_str = None
+
+    if isinstance(type_obj, str):
+        type_str = type_obj
+    elif hasattr(type_obj, "__name__"):
+        type_str = type_obj.__name__
+    elif hasattr(type_obj, "name"):
+        type_str = type_obj.name
+    else:
+        raise RuntimeError("Cannot infer type name from input")
+
+    assert type_str in _str_to_ctype.keys()
+
+    my_dtype = np.dtype(type_str)
+    my_ctype = _str_to_ctype[type_str]
+
+    assert my_dtype.itemsize == ctypes.sizeof(my_ctype)
+
+    return my_dtype, my_ctype
+
+
+def is_pickleable(obj: Any) -> bool:
+    try:
+        with io.BytesIO() as stream:
+            pickle.dump(obj, stream)
+        return True
+    except:
+        return False
+
+
+# Functionality to import modules/objects by name, and call functions by name
+# ------------------------------------------------------------------------------------------
+
+def get_module_from_obj_name(obj_name: str) -> Tuple[types.ModuleType, str]:
+    """Searches for the underlying module behind the name to some python object.
+    Returns the module and the object name (original name with module part removed)."""
+
+    # allow convenience shorthands, substitute them by full names
+    obj_name = re.sub("^np.", "numpy.", obj_name)
+    obj_name = re.sub("^tf.", "tensorflow.", obj_name)
+
+    # list alternatives for (module_name, local_obj_name)
+    parts = obj_name.split(".")
+    name_pairs = [(".".join(parts[:i]), ".".join(parts[i:])) for i in range(len(parts), 0, -1)]
+
+    # try each alternative in turn
+    for module_name, local_obj_name in name_pairs:
+        try:
+            module = importlib.import_module(module_name) # may raise ImportError
+            get_obj_from_module(module, local_obj_name) # may raise AttributeError
+            return module, local_obj_name
+        except:
+            pass
+
+    # maybe some of the modules themselves contain errors?
+    for module_name, _local_obj_name in name_pairs:
+        try:
+            importlib.import_module(module_name) # may raise ImportError
+        except ImportError:
+            if not str(sys.exc_info()[1]).startswith("No module named '" + module_name + "'"):
+                raise
+
+    # maybe the requested attribute is missing?
+    for module_name, local_obj_name in name_pairs:
+        try:
+            module = importlib.import_module(module_name) # may raise ImportError
+            get_obj_from_module(module, local_obj_name) # may raise AttributeError
+        except ImportError:
+            pass
+
+    # we are out of luck, but we have no idea why
+    raise ImportError(obj_name)
+
+
+def get_obj_from_module(module: types.ModuleType, obj_name: str) -> Any:
+    """Traverses the object name and returns the last (rightmost) python object."""
+    if obj_name == '':
+        return module
+    obj = module
+    for part in obj_name.split("."):
+        obj = getattr(obj, part)
+    return obj
+
+
+def get_obj_by_name(name: str) -> Any:
+    """Finds the python object with the given name."""
+    module, obj_name = get_module_from_obj_name(name)
+    return get_obj_from_module(module, obj_name)
+
+
+def call_func_by_name(*args, func_name: str = None, **kwargs) -> Any:
+    """Finds the python object with the given name and calls it as a function."""
+    assert func_name is not None
+    func_obj = get_obj_by_name(func_name)
+    assert callable(func_obj)
+    return func_obj(*args, **kwargs)
+
+
+def get_module_dir_by_obj_name(obj_name: str) -> str:
+    """Get the directory path of the module containing the given object name."""
+    module, _ = get_module_from_obj_name(obj_name)
+    return os.path.dirname(inspect.getfile(module))
+
+
+def is_top_level_function(obj: Any) -> bool:
+    """Determine whether the given object is a top-level function, i.e., defined at module scope using 'def'."""
+    return callable(obj) and obj.__name__ in sys.modules[obj.__module__].__dict__
+
+
+def get_top_level_function_name(obj: Any) -> str:
+    """Return the fully-qualified name of a top-level function."""
+    assert is_top_level_function(obj)
+    return obj.__module__ + "." + obj.__name__
+
+
+# File system helpers
+# ------------------------------------------------------------------------------------------
+
+def list_dir_recursively_with_ignore(dir_path: str, ignores: List[str] = None, add_base_to_relative: bool = False) -> List[Tuple[str, str]]:
+    """List all files recursively in a given directory while ignoring given file and directory names.
+    Returns list of tuples containing both absolute and relative paths."""
+    assert os.path.isdir(dir_path)
+    base_name = os.path.basename(os.path.normpath(dir_path))
+
+    if ignores is None:
+        ignores = []
+
+    result = []
+
+    for root, dirs, files in os.walk(dir_path, topdown=True):
+        for ignore_ in ignores:
+            dirs_to_remove = [d for d in dirs if fnmatch.fnmatch(d, ignore_)]
+
+            # dirs need to be edited in-place
+            for d in dirs_to_remove:
+                dirs.remove(d)
+
+            files = [f for f in files if not fnmatch.fnmatch(f, ignore_)]
+
+        absolute_paths = [os.path.join(root, f) for f in files]
+        relative_paths = [os.path.relpath(p, dir_path) for p in absolute_paths]
+
+        if add_base_to_relative:
+            relative_paths = [os.path.join(base_name, p) for p in relative_paths]
+
+        assert len(absolute_paths) == len(relative_paths)
+        result += zip(absolute_paths, relative_paths)
+
+    return result
+
+
+def copy_files_and_create_dirs(files: List[Tuple[str, str]]) -> None:
+    """Takes in a list of tuples of (src, dst) paths and copies files.
+    Will create all necessary directories."""
+    for file in files:
+        target_dir_name = os.path.dirname(file[1])
+
+        # will create all intermediate-level directories
+        if not os.path.exists(target_dir_name):
+            os.makedirs(target_dir_name)
+
+        shutil.copyfile(file[0], file[1])
+
+
+# URL helpers
+# ------------------------------------------------------------------------------------------
+
+def is_url(obj: Any) -> bool:
+    """Determine whether the given object is a valid URL string."""
+    if not isinstance(obj, str) or not "://" in obj:
+        return False
+    try:
+        res = requests.compat.urlparse(obj)
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+        res = requests.compat.urlparse(requests.compat.urljoin(obj, "/"))
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+    except:
+        return False
+    return True
+
+
+def open_url(url: str, cache_dir: str = None, num_attempts: int = 10, verbose: bool = True) -> Any:
+    """Download the given URL and return a binary-mode file object to access the data."""
+    if not is_url(url) and os.path.isfile(url):
+        return open(url, 'rb')
+
+    assert is_url(url)
+    assert num_attempts >= 1
+
+    # Lookup from cache.
+    url_md5 = hashlib.md5(url.encode("utf-8")).hexdigest()
+    if cache_dir is not None:
+        cache_files = glob.glob(os.path.join(cache_dir, url_md5 + "_*"))
+        if len(cache_files) == 1:
+            return open(cache_files[0], "rb")
+
+    # Download.
+    url_name = None
+    url_data = None
+    with requests.Session() as session:
+        if verbose:
+            print("Downloading %s ..." % url, end="", flush=True)
+        for attempts_left in reversed(range(num_attempts)):
+            try:
+                with session.get(url) as res:
+                    res.raise_for_status()
+                    if len(res.content) == 0:
+                        raise IOError("No data received")
+
+                    if len(res.content) < 8192:
+                        content_str = res.content.decode("utf-8")
+                        if "download_warning" in res.headers.get("Set-Cookie", ""):
+                            links = [html.unescape(link) for link in content_str.split('"') if "export=download" in link]
+                            if len(links) == 1:
+                                url = requests.compat.urljoin(url, links[0])
+                                raise IOError("Google Drive virus checker nag")
+                        if "Google Drive - Quota exceeded" in content_str:
+                            raise IOError("Google Drive quota exceeded")
+
+                    match = re.search(r'filename="([^"]*)"', res.headers.get("Content-Disposition", ""))
+                    url_name = match[1] if match else url
+                    url_data = res.content
+                    if verbose:
+                        print(" done")
+                    break
+            except:
+                if not attempts_left:
+                    if verbose:
+                        print(" failed")
+                    raise
+                if verbose:
+                    print(".", end="", flush=True)
+
+    # Save to cache.
+    if cache_dir is not None:
+        safe_name = re.sub(r"[^0-9a-zA-Z-._]", "_", url_name)
+        cache_file = os.path.join(cache_dir, url_md5 + "_" + safe_name)
+        temp_file = os.path.join(cache_dir, "tmp_" + uuid.uuid4().hex + "_" + url_md5 + "_" + safe_name)
+        os.makedirs(cache_dir, exist_ok=True)
+        with open(temp_file, "wb") as f:
+            f.write(url_data)
+        os.replace(temp_file, cache_file) # atomic
+
+    # Return data as file object.
+    return io.BytesIO(url_data)

encode_images.py

@@ -0,0 +1,242 @@
+import os
+import argparse
+import pickle
+from tqdm import tqdm
+import PIL.Image
+from PIL import ImageFilter
+import numpy as np
+import dnnlib
+import dnnlib.tflib as tflib
+import config
+from encoder.generator_model import Generator
+from encoder.perceptual_model import PerceptualModel, load_images
+#from tensorflow.keras.models import load_model
+from keras.models import load_model
+from keras.applications.resnet50 import preprocess_input
+
+def split_to_batches(l, n):
+    for i in range(0, len(l), n):
+        yield l[i:i + n]
+
+def str2bool(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 main():
+    parser = argparse.ArgumentParser(description='Find latent representation of reference images using perceptual losses', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument('src_dir', help='Directory with images for encoding')
+    parser.add_argument('generated_images_dir', help='Directory for storing generated images')
+    parser.add_argument('dlatent_dir', help='Directory for storing dlatent representations')
+    parser.add_argument('--data_dir', default='data', help='Directory for storing optional models')
+    parser.add_argument('--mask_dir', default='masks', help='Directory for storing optional masks')
+    parser.add_argument('--load_last', default='', help='Start with embeddings from directory')
+    parser.add_argument('--dlatent_avg', default='', help='Use dlatent from file specified here for truncation instead of dlatent_avg from Gs')
+    parser.add_argument('--model_url', default='./data/karras2019stylegan-ffhq-1024x1024.pkl', help='Fetch a StyleGAN model to train on from this URL')
+    parser.add_argument('--architecture', default='./data/vgg16_zhang_perceptual.pkl', help='Сonvolutional neural network model from this URL')
+    parser.add_argument('--model_res', default=1024, help='The dimension of images in the StyleGAN model', type=int)
+    parser.add_argument('--batch_size', default=1, help='Batch size for generator and perceptual model', type=int)
+    parser.add_argument('--optimizer', default='ggt', help='Optimization algorithm used for optimizing dlatents')
+
+    # Perceptual model params
+    parser.add_argument('--image_size', default=256, help='Size of images for perceptual model', type=int)
+    parser.add_argument('--resnet_image_size', default=256, help='Size of images for the Resnet model', type=int)
+    parser.add_argument('--lr', default=0.25, help='Learning rate for perceptual model', type=float)
+    parser.add_argument('--decay_rate', default=0.9, help='Decay rate for learning rate', type=float)
+    parser.add_argument('--iterations', default=100, help='Number of optimization steps for each batch', type=int)
+    parser.add_argument('--decay_steps', default=4, help='Decay steps for learning rate decay (as a percent of iterations)', type=float)
+    parser.add_argument('--early_stopping', default=True, help='Stop early once training stabilizes', type=str2bool, nargs='?', const=True)
+    parser.add_argument('--early_stopping_threshold', default=0.5, help='Stop after this threshold has been reached', type=float)
+    parser.add_argument('--early_stopping_patience', default=10, help='Number of iterations to wait below threshold', type=int)    
+    parser.add_argument('--load_effnet', default='data/finetuned_effnet.h5', help='Model to load for EfficientNet approximation of dlatents')
+    parser.add_argument('--load_resnet', default='data/finetuned_resnet.h5', help='Model to load for ResNet approximation of dlatents')
+    parser.add_argument('--use_preprocess_input', default=True, help='Call process_input() first before using feed forward net', type=str2bool, nargs='?', const=True)
+    parser.add_argument('--use_best_loss', default=True, help='Output the lowest loss value found as the solution', type=str2bool, nargs='?', const=True)
+    parser.add_argument('--average_best_loss', default=0.25, help='Do a running weighted average with the previous best dlatents found', type=float)
+    parser.add_argument('--sharpen_input', default=True, help='Sharpen the input images', type=str2bool, nargs='?', const=True)
+
+    # Loss function options
+    parser.add_argument('--use_vgg_loss', default=0.4, help='Use VGG perceptual loss; 0 to disable, > 0 to scale.', type=float)
+    parser.add_argument('--use_vgg_layer', default=9, help='Pick which VGG layer to use.', type=int)
+    parser.add_argument('--use_pixel_loss', default=1.5, help='Use logcosh image pixel loss; 0 to disable, > 0 to scale.', type=float)
+    parser.add_argument('--use_mssim_loss', default=200, help='Use MS-SIM perceptual loss; 0 to disable, > 0 to scale.', type=float)
+    parser.add_argument('--use_lpips_loss', default=100, help='Use LPIPS perceptual loss; 0 to disable, > 0 to scale.', type=float)
+    parser.add_argument('--use_l1_penalty', default=0.5, help='Use L1 penalty on latents; 0 to disable, > 0 to scale.', type=float)
+    parser.add_argument('--use_discriminator_loss', default=0.5, help='Use trained discriminator to evaluate realism.', type=float)
+    parser.add_argument('--use_adaptive_loss', default=False, help='Use the adaptive robust loss function from Google Research for pixel and VGG feature loss.', type=str2bool, nargs='?', const=True)
+
+    # Generator params
+    parser.add_argument('--randomize_noise', default=False, help='Add noise to dlatents during optimization', type=str2bool, nargs='?', const=True)
+    parser.add_argument('--tile_dlatents', default=False, help='Tile dlatents to use a single vector at each scale', type=str2bool, nargs='?', const=True)
+    parser.add_argument('--clipping_threshold', default=2.0, help='Stochastic clipping of gradient values outside of this threshold', type=float)
+
+    # Masking params
+    parser.add_argument('--load_mask', default=False, help='Load segmentation masks', type=str2bool, nargs='?', const=True)
+    parser.add_argument('--face_mask', default=True, help='Generate a mask for predicting only the face area', type=str2bool, nargs='?', const=True)
+    parser.add_argument('--use_grabcut', default=True, help='Use grabcut algorithm on the face mask to better segment the foreground', type=str2bool, nargs='?', const=True)
+    parser.add_argument('--scale_mask', default=1.4, help='Look over a wider section of foreground for grabcut', type=float)
+    parser.add_argument('--composite_mask', default=True, help='Merge the unmasked area back into the generated image', type=str2bool, nargs='?', const=True)
+    parser.add_argument('--composite_blur', default=8, help='Size of blur filter to smoothly composite the images', type=int)
+
+    # Video params
+    parser.add_argument('--video_dir', default='videos', help='Directory for storing training videos')
+    parser.add_argument('--output_video', default=False, help='Generate videos of the optimization process', type=bool)
+    parser.add_argument('--video_codec', default='MJPG', help='FOURCC-supported video codec name')
+    parser.add_argument('--video_frame_rate', default=24, help='Video frames per second', type=int)
+    parser.add_argument('--video_size', default=512, help='Video size in pixels', type=int)
+    parser.add_argument('--video_skip', default=1, help='Only write every n frames (1 = write every frame)', type=int)
+
+    args, other_args = parser.parse_known_args()
+
+    args.decay_steps *= 0.01 * args.iterations # Calculate steps as a percent of total iterations
+
+    if args.output_video:
+      import cv2
+      synthesis_kwargs = dict(output_transform=dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=False), minibatch_size=args.batch_size)
+
+    ref_images = [os.path.join(args.src_dir, x) for x in os.listdir(args.src_dir)]
+    ref_images = list(filter(os.path.isfile, ref_images))
+
+    if len(ref_images) == 0:
+        raise Exception('%s is empty' % args.src_dir)
+
+    os.makedirs(args.data_dir, exist_ok=True)
+    os.makedirs(args.mask_dir, exist_ok=True)
+    os.makedirs(args.generated_images_dir, exist_ok=True)
+    os.makedirs(args.dlatent_dir, exist_ok=True)
+    os.makedirs(args.video_dir, exist_ok=True)
+
+    # Initialize generator and perceptual model
+    tflib.init_tf()
+    with dnnlib.util.open_url(args.model_url, cache_dir=config.cache_dir) as f:
+        generator_network, discriminator_network, Gs_network = pickle.load(f)
+
+    generator = Generator(Gs_network, args.batch_size, clipping_threshold=args.clipping_threshold, tiled_dlatent=args.tile_dlatents, model_res=args.model_res, randomize_noise=args.randomize_noise)
+    if (args.dlatent_avg != ''):
+        generator.set_dlatent_avg(np.load(args.dlatent_avg))
+
+    perc_model = None
+    if (args.use_lpips_loss > 0.00000001):
+        with dnnlib.util.open_url(args.architecture, cache_dir=config.cache_dir) as f:
+            perc_model =  pickle.load(f)
+    perceptual_model = PerceptualModel(args, perc_model=perc_model, batch_size=args.batch_size)
+    perceptual_model.build_perceptual_model(generator, discriminator_network)
+
+    ff_model = None
+
+    # Optimize (only) dlatents by minimizing perceptual loss between reference and generated images in feature space
+    for images_batch in tqdm(split_to_batches(ref_images, args.batch_size), total=len(ref_images)//args.batch_size):
+        names = [os.path.splitext(os.path.basename(x))[0] for x in images_batch]
+        if args.output_video:
+          video_out = {}
+          for name in names:
+            video_out[name] = cv2.VideoWriter(os.path.join(args.video_dir, f'{name}.avi'),cv2.VideoWriter_fourcc(*args.video_codec), args.video_frame_rate, (args.video_size,args.video_size))
+
+        perceptual_model.set_reference_images(images_batch)
+        dlatents = None
+        if (args.load_last != ''): # load previous dlatents for initialization
+            for name in names:
+                dl = np.expand_dims(np.load(os.path.join(args.load_last, f'{name}.npy')),axis=0)
+                if (dlatents is None):
+                    dlatents = dl
+                else:
+                    dlatents = np.vstack((dlatents,dl))
+        else:
+            if (ff_model is None):
+                if os.path.exists(args.load_resnet):
+                    from keras.applications.resnet50 import preprocess_input
+                    print("Loading ResNet Model:")
+                    ff_model = load_model(args.load_resnet)
+            if (ff_model is None):
+                if os.path.exists(args.load_effnet):
+                    import efficientnet
+                    from efficientnet import preprocess_input
+                    print("Loading EfficientNet Model:")
+                    ff_model = load_model(args.load_effnet)
+            if (ff_model is not None): # predict initial dlatents with ResNet model
+                if (args.use_preprocess_input):
+                    dlatents = ff_model.predict(preprocess_input(load_images(images_batch,image_size=args.resnet_image_size)))
+                else:
+                    dlatents = ff_model.predict(load_images(images_batch,image_size=args.resnet_image_size))
+        if dlatents is not None:
+            generator.set_dlatents(dlatents)
+        op = perceptual_model.optimize(generator.dlatent_variable, iterations=args.iterations, use_optimizer=args.optimizer)
+        pbar = tqdm(op, leave=False, total=args.iterations)
+        vid_count = 0
+        best_loss = None
+        best_dlatent = None
+        avg_loss_count = 0
+        if args.early_stopping:
+            avg_loss = prev_loss = None
+        for loss_dict in pbar:
+            if args.early_stopping: # early stopping feature
+                if prev_loss is not None:
+                    if avg_loss is not None:
+                        avg_loss = 0.5 * avg_loss + (prev_loss - loss_dict["loss"])
+                        if avg_loss < args.early_stopping_threshold: # count while under threshold; else reset
+                            avg_loss_count += 1
+                        else:
+                            avg_loss_count = 0
+                        if avg_loss_count > args.early_stopping_patience: # stop once threshold is reached
+                            print("")
+                            break
+                    else:
+                        avg_loss = prev_loss - loss_dict["loss"]
+            pbar.set_description(" ".join(names) + ": " + "; ".join(["{} {:.4f}".format(k, v) for k, v in loss_dict.items()]))
+            if best_loss is None or loss_dict["loss"] < best_loss:
+                if best_dlatent is None or args.average_best_loss <= 0.00000001:
+                    best_dlatent = generator.get_dlatents()
+                else:
+                    best_dlatent = 0.25 * best_dlatent + 0.75 * generator.get_dlatents()
+                if args.use_best_loss:
+                    generator.set_dlatents(best_dlatent)
+                best_loss = loss_dict["loss"]
+            if args.output_video and (vid_count % args.video_skip == 0):
+              batch_frames = generator.generate_images()
+              for i, name in enumerate(names):
+                video_frame = PIL.Image.fromarray(batch_frames[i], 'RGB').resize((args.video_size,args.video_size),PIL.Image.LANCZOS)
+                video_out[name].write(cv2.cvtColor(np.array(video_frame).astype('uint8'), cv2.COLOR_RGB2BGR))
+            generator.stochastic_clip_dlatents()
+            prev_loss = loss_dict["loss"]
+        if not args.use_best_loss:
+            best_loss = prev_loss
+        print(" ".join(names), " Loss {:.4f}".format(best_loss))
+
+        if args.output_video:
+            for name in names:
+                video_out[name].release()
+
+        # Generate images from found dlatents and save them
+        if args.use_best_loss:
+            generator.set_dlatents(best_dlatent)
+        generated_images = generator.generate_images()
+        generated_dlatents = generator.get_dlatents()
+        for img_array, dlatent, img_path, img_name in zip(generated_images, generated_dlatents, images_batch, names):
+            mask_img = None
+            if args.composite_mask and (args.load_mask or args.face_mask):
+                _, im_name = os.path.split(img_path)
+                mask_img = os.path.join(args.mask_dir, f'{im_name}')
+            if args.composite_mask and mask_img is not None and os.path.isfile(mask_img):
+                orig_img = PIL.Image.open(img_path).convert('RGB')
+                width, height = orig_img.size
+                imask = PIL.Image.open(mask_img).convert('L').resize((width, height))
+                imask = imask.filter(ImageFilter.GaussianBlur(args.composite_blur))
+                mask = np.array(imask)/255
+                mask = np.expand_dims(mask,axis=-1)
+                img_array = mask*np.array(img_array) + (1.0-mask)*np.array(orig_img)
+                img_array = img_array.astype(np.uint8)
+                #img_array = np.where(mask, np.array(img_array), orig_img)
+            img = PIL.Image.fromarray(img_array, 'RGB')
+            img.save(os.path.join(args.generated_images_dir, f'{img_name}.png'), 'PNG')
+            np.save(os.path.join(args.dlatent_dir, f'{img_name}.npy'), dlatent)
+
+        generator.reset_dlatents()
+
+
+if __name__ == "__main__":
+    main()

encoder/generator_model.py

@@ -0,0 +1,137 @@
+import math
+import tensorflow as tf
+import numpy as np
+import dnnlib.tflib as tflib
+from functools import partial
+
+
+def create_stub(name, batch_size):
+    return tf.constant(0, dtype='float32', shape=(batch_size, 0))
+
+
+def create_variable_for_generator(name, batch_size, tiled_dlatent, model_scale=18, tile_size = 1):
+    if tiled_dlatent:
+        low_dim_dlatent = tf.get_variable('learnable_dlatents',
+            shape=(batch_size, tile_size, 512),
+            dtype='float32',
+            initializer=tf.initializers.random_normal())
+        return tf.tile(low_dim_dlatent, [1, model_scale // tile_size, 1])
+    else:
+        return tf.get_variable('learnable_dlatents',
+            shape=(batch_size, model_scale, 512),
+            dtype='float32',
+            initializer=tf.initializers.random_normal())
+
+
+class Generator:
+    def __init__(self, model, batch_size, custom_input=None, clipping_threshold=2, tiled_dlatent=False, model_res=1024, randomize_noise=False):
+        self.batch_size = batch_size
+        self.tiled_dlatent=tiled_dlatent
+        self.model_scale = int(2*(math.log(model_res,2)-1)) # For example, 1024 -> 18
+
+        if tiled_dlatent:
+            self.initial_dlatents = np.zeros((self.batch_size, 512))
+            model.components.synthesis.run(np.zeros((self.batch_size, self.model_scale, 512)),
+                randomize_noise=randomize_noise, minibatch_size=self.batch_size,
+                custom_inputs=[partial(create_variable_for_generator, batch_size=batch_size, tiled_dlatent=True),
+                                                partial(create_stub, batch_size=batch_size)],
+                structure='fixed')
+        else:
+            self.initial_dlatents = np.zeros((self.batch_size, self.model_scale, 512))
+            if custom_input is not None:
+                model.components.synthesis.run(self.initial_dlatents,
+                    randomize_noise=randomize_noise, minibatch_size=self.batch_size,
+                    custom_inputs=[partial(custom_input.eval(), batch_size=batch_size), partial(create_stub, batch_size=batch_size)],
+                    structure='fixed')
+            else:
+                model.components.synthesis.run(self.initial_dlatents,
+                    randomize_noise=randomize_noise, minibatch_size=self.batch_size,
+                    custom_inputs=[partial(create_variable_for_generator, batch_size=batch_size, tiled_dlatent=False, model_scale=self.model_scale),
+                                                    partial(create_stub, batch_size=batch_size)],
+                    structure='fixed')
+
+        self.dlatent_avg_def = model.get_var('dlatent_avg')
+        self.reset_dlatent_avg()
+        self.sess = tf.compat.v1.get_default_session()
+        self.graph = tf.compat.v1.get_default_graph()
+
+        self.dlatent_variable = next(v for v in tf.compat.v1.global_variables() if 'learnable_dlatents' in v.name)
+        self._assign_dlatent_ph = tf.compat.v1.placeholder(tf.float32, name="assign_dlatent_ph")
+        self._assign_dlantent = tf.assign(self.dlatent_variable, self._assign_dlatent_ph)
+        self.set_dlatents(self.initial_dlatents)
+
+        def get_tensor(name):
+            try:
+                return self.graph.get_tensor_by_name(name)
+            except KeyError:
+                return None
+
+        self.generator_output = get_tensor('G_synthesis_1/_Run/concat:0')
+        if self.generator_output is None:
+            self.generator_output = get_tensor('G_synthesis_1/_Run/concat/concat:0')
+        if self.generator_output is None:
+            self.generator_output = get_tensor('G_synthesis_1/_Run/concat_1/concat:0')
+        # If we loaded only Gs and didn't load G or D, then scope "G_synthesis_1" won't exist in the graph.
+        if self.generator_output is None:
+            self.generator_output = get_tensor('G_synthesis/_Run/concat:0')
+        if self.generator_output is None:
+            self.generator_output = get_tensor('G_synthesis/_Run/concat/concat:0')
+        if self.generator_output is None:
+            self.generator_output = get_tensor('G_synthesis/_Run/concat_1/concat:0')
+        if self.generator_output is None:
+            for op in self.graph.get_operations():
+                print(op)
+            raise Exception("Couldn't find G_synthesis_1/_Run/concat tensor output")
+        self.generated_image = tflib.convert_images_to_uint8(self.generator_output, nchw_to_nhwc=True, uint8_cast=False)
+        self.generated_image_uint8 = tf.saturate_cast(self.generated_image, tf.uint8)
+
+        # Implement stochastic clipping similar to what is described in https://arxiv.org/abs/1702.04782
+        # (Slightly different in that the latent space is normal gaussian here and was uniform in [-1, 1] in that paper,
+        # so we clip any vector components outside of [-2, 2]. It seems fine, but I haven't done an ablation check.)
+        clipping_mask = tf.math.logical_or(self.dlatent_variable > clipping_threshold, self.dlatent_variable < -clipping_threshold)
+        clipped_values = tf.where(clipping_mask, tf.random.normal(shape=self.dlatent_variable.shape), self.dlatent_variable)
+        self.stochastic_clip_op = tf.assign(self.dlatent_variable, clipped_values)
+
+    def reset_dlatents(self):
+        self.set_dlatents(self.initial_dlatents)
+
+    def set_dlatents(self, dlatents):
+        if self.tiled_dlatent:
+            if (dlatents.shape != (self.batch_size, 512)) and (dlatents.shape[1] != 512):
+                dlatents = np.mean(dlatents, axis=1)
+            if (dlatents.shape != (self.batch_size, 512)):
+                dlatents = np.vstack([dlatents, np.zeros((self.batch_size-dlatents.shape[0], 512))])
+            assert (dlatents.shape == (self.batch_size, 512))
+        else:
+            if (dlatents.shape[1] > self.model_scale):
+                dlatents = dlatents[:,:self.model_scale,:]
+            if (isinstance(dlatents.shape[0], int)):
+                if (dlatents.shape != (self.batch_size, self.model_scale, 512)):
+                    dlatents = np.vstack([dlatents, np.zeros((self.batch_size-dlatents.shape[0], self.model_scale, 512))])
+                assert (dlatents.shape == (self.batch_size, self.model_scale, 512))
+                self.sess.run([self._assign_dlantent], {self._assign_dlatent_ph: dlatents})
+                return
+            else:
+                self._assign_dlantent = tf.assign(self.dlatent_variable, dlatents)
+                return
+        self.sess.run([self._assign_dlantent], {self._assign_dlatent_ph: dlatents})
+
+    def stochastic_clip_dlatents(self):
+        self.sess.run(self.stochastic_clip_op)
+
+    def get_dlatents(self):
+        return self.sess.run(self.dlatent_variable)
+
+    def get_dlatent_avg(self):
+        return self.dlatent_avg
+
+    def set_dlatent_avg(self, dlatent_avg):
+        self.dlatent_avg = dlatent_avg
+
+    def reset_dlatent_avg(self):
+        self.dlatent_avg = self.dlatent_avg_def
+
+    def generate_images(self, dlatents=None):
+        if dlatents is not None:
+            self.set_dlatents(dlatents)
+        return self.sess.run(self.generated_image_uint8)

encoder/perceptual_model.py

@@ -0,0 +1,304 @@
+from __future__ import absolute_import, division, print_function, unicode_literals
+import tensorflow as tf
+#import tensorflow_probability as tfp
+#tf.enable_eager_execution()
+
+import os
+import bz2
+import PIL.Image
+from PIL import ImageFilter
+import numpy as np
+from keras.models import Model
+from keras.utils import get_file
+from keras.applications.vgg16 import VGG16, preprocess_input
+import keras.backend as K
+import traceback
+import dnnlib.tflib as tflib
+
+def load_images(images_list, image_size=256, sharpen=False):
+    loaded_images = list()
+    for img_path in images_list:
+      img = PIL.Image.open(img_path).convert('RGB')
+      if image_size is not None:
+        img = img.resize((image_size,image_size),PIL.Image.LANCZOS)
+        if (sharpen):
+            img = img.filter(ImageFilter.DETAIL)
+      img = np.array(img)
+      img = np.expand_dims(img, 0)
+      loaded_images.append(img)
+    loaded_images = np.vstack(loaded_images)
+    return loaded_images
+
+def tf_custom_adaptive_loss(a,b):
+    from adaptive import lossfun
+    shape = a.get_shape().as_list()
+    dim = np.prod(shape[1:])
+    a = tf.reshape(a, [-1, dim])
+    b = tf.reshape(b, [-1, dim])
+    loss, _, _ = lossfun(b-a, var_suffix='1')
+    return tf.math.reduce_mean(loss)
+
+def tf_custom_adaptive_rgb_loss(a,b):
+    from adaptive import image_lossfun
+    loss, _, _ = image_lossfun(b-a, color_space='RGB', representation='PIXEL')
+    return tf.math.reduce_mean(loss)
+
+def tf_custom_l1_loss(img1,img2):
+  return tf.math.reduce_mean(tf.math.abs(img2-img1), axis=None)
+
+def tf_custom_logcosh_loss(img1,img2):
+  return tf.math.reduce_mean(tf.keras.losses.logcosh(img1,img2))
+
+def create_stub(batch_size):
+    return tf.constant(0, dtype='float32', shape=(batch_size, 0))
+
+def unpack_bz2(src_path):
+    data = bz2.BZ2File(src_path).read()
+    dst_path = src_path[:-4]
+    with open(dst_path, 'wb') as fp:
+        fp.write(data)
+    return dst_path
+
+class PerceptualModel:
+    def __init__(self, args, batch_size=1, perc_model=None, sess=None):
+        self.sess = tf.compat.v1.get_default_session() if sess is None else sess
+        K.set_session(self.sess)
+        self.epsilon = 0.00000001
+        self.lr = args.lr
+        self.decay_rate = args.decay_rate
+        self.decay_steps = args.decay_steps
+        self.img_size = args.image_size
+        self.layer = args.use_vgg_layer
+        self.vgg_loss = args.use_vgg_loss
+        self.face_mask = args.face_mask
+        self.use_grabcut = args.use_grabcut
+        self.scale_mask = args.scale_mask
+        self.mask_dir = args.mask_dir
+        if (self.layer <= 0 or self.vgg_loss <= self.epsilon):
+            self.vgg_loss = None
+        self.pixel_loss = args.use_pixel_loss
+        if (self.pixel_loss <= self.epsilon):
+            self.pixel_loss = None
+        self.mssim_loss = args.use_mssim_loss
+        if (self.mssim_loss <= self.epsilon):
+            self.mssim_loss = None
+        self.lpips_loss = args.use_lpips_loss
+        if (self.lpips_loss <= self.epsilon):
+            self.lpips_loss = None
+        self.l1_penalty = args.use_l1_penalty
+        if (self.l1_penalty <= self.epsilon):
+            self.l1_penalty = None
+        self.adaptive_loss = args.use_adaptive_loss
+        self.sharpen_input = args.sharpen_input
+        self.batch_size = batch_size
+        if perc_model is not None and self.lpips_loss is not None:
+            self.perc_model = perc_model
+        else:
+            self.perc_model = None
+        self.ref_img = None
+        self.ref_weight = None
+        self.perceptual_model = None
+        self.ref_img_features = None
+        self.features_weight = None
+        self.loss = None
+        self.discriminator_loss = args.use_discriminator_loss
+        if (self.discriminator_loss <= self.epsilon):
+            self.discriminator_loss = None
+        if self.discriminator_loss is not None:
+            self.discriminator = None
+            self.stub = create_stub(batch_size)
+
+        if self.face_mask:
+            import dlib
+            self.detector = dlib.get_frontal_face_detector()
+            landmarks_model_path = unpack_bz2('shape_predictor_68_face_landmarks.dat.bz2')
+            self.predictor = dlib.shape_predictor(landmarks_model_path)
+
+    def add_placeholder(self, var_name):
+        var_val = getattr(self, var_name)
+        setattr(self, var_name + "_placeholder", tf.compat.v1.placeholder(var_val.dtype, shape=var_val.get_shape()))
+        setattr(self, var_name + "_op", var_val.assign(getattr(self, var_name + "_placeholder")))
+
+    def assign_placeholder(self, var_name, var_val):
+        self.sess.run(getattr(self, var_name + "_op"), {getattr(self, var_name + "_placeholder"): var_val})
+
+    def build_perceptual_model(self, generator, discriminator=None):
+        # Learning rate
+        global_step = tf.Variable(0, dtype=tf.int32, trainable=False, name="global_step")
+        incremented_global_step = tf.compat.v1.assign_add(global_step, 1)
+        self._reset_global_step = tf.assign(global_step, 0)
+        self.learning_rate = tf.compat.v1.train.exponential_decay(self.lr, incremented_global_step,
+                self.decay_steps, self.decay_rate, staircase=True)
+        self.sess.run([self._reset_global_step])
+
+        if self.discriminator_loss is not None:
+            self.discriminator = discriminator
+
+        generated_image_tensor = generator.generated_image
+        generated_image = tf.compat.v1.image.resize_nearest_neighbor(generated_image_tensor,
+                                                                  (self.img_size, self.img_size), align_corners=True)
+
+        self.ref_img = tf.get_variable('ref_img', shape=generated_image.shape,
+                                                dtype='float32', initializer=tf.initializers.zeros())
+        self.ref_weight = tf.get_variable('ref_weight', shape=generated_image.shape,
+                                               dtype='float32', initializer=tf.initializers.zeros())
+        self.add_placeholder("ref_img")
+        self.add_placeholder("ref_weight")
+
+        if (self.vgg_loss is not None):
+            vgg16 = VGG16(include_top=False, weights='vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5', input_shape=(self.img_size, self.img_size, 3)) # https://github.com/fchollet/deep-learning-models/releases/download/v0.1/vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5
+            self.perceptual_model = Model(vgg16.input, vgg16.layers[self.layer].output)
+            generated_img_features = self.perceptual_model(preprocess_input(self.ref_weight * generated_image))
+            self.ref_img_features = tf.get_variable('ref_img_features', shape=generated_img_features.shape,
+                                                dtype='float32', initializer=tf.initializers.zeros())
+            self.features_weight = tf.get_variable('features_weight', shape=generated_img_features.shape,
+                                               dtype='float32', initializer=tf.initializers.zeros())
+            self.sess.run([self.features_weight.initializer, self.features_weight.initializer])
+            self.add_placeholder("ref_img_features")
+            self.add_placeholder("features_weight")
+
+        if self.perc_model is not None and self.lpips_loss is not None:
+            img1 = tflib.convert_images_from_uint8(self.ref_weight * self.ref_img, nhwc_to_nchw=True)
+            img2 = tflib.convert_images_from_uint8(self.ref_weight * generated_image, nhwc_to_nchw=True)
+
+        self.loss = 0
+        # L1 loss on VGG16 features
+        if (self.vgg_loss is not None):
+            if self.adaptive_loss:
+                self.loss += self.vgg_loss * tf_custom_adaptive_loss(self.features_weight * self.ref_img_features, self.features_weight * generated_img_features)
+            else:
+                self.loss += self.vgg_loss * tf_custom_logcosh_loss(self.features_weight * self.ref_img_features, self.features_weight * generated_img_features)
+        # + logcosh loss on image pixels
+        if (self.pixel_loss is not None):
+            if self.adaptive_loss:
+                self.loss += self.pixel_loss * tf_custom_adaptive_rgb_loss(self.ref_weight * self.ref_img, self.ref_weight * generated_image)
+            else:
+                self.loss += self.pixel_loss * tf_custom_logcosh_loss(self.ref_weight * self.ref_img, self.ref_weight * generated_image)
+        # + MS-SIM loss on image pixels
+        if (self.mssim_loss is not None):
+            self.loss += self.mssim_loss * tf.math.reduce_mean(1-tf.image.ssim_multiscale(self.ref_weight * self.ref_img, self.ref_weight * generated_image, 1))
+        # + extra perceptual loss on image pixels
+        if self.perc_model is not None and self.lpips_loss is not None:
+            self.loss += self.lpips_loss * tf.math.reduce_mean(self.perc_model.get_output_for(img1, img2))
+        # + L1 penalty on dlatent weights
+        if self.l1_penalty is not None:
+            self.loss += self.l1_penalty * 512 * tf.math.reduce_mean(tf.math.abs(generator.dlatent_variable-generator.get_dlatent_avg()))
+        # discriminator loss (realism)
+        if self.discriminator_loss is not None:
+            self.loss += self.discriminator_loss * tf.math.reduce_mean(self.discriminator.get_output_for(tflib.convert_images_from_uint8(generated_image_tensor, nhwc_to_nchw=True), self.stub))
+        # - discriminator_network.get_output_for(tflib.convert_images_from_uint8(ref_img, nhwc_to_nchw=True), stub)
+
+
+    def generate_face_mask(self, im):
+        from imutils import face_utils
+        import cv2
+        rects = self.detector(im, 1)
+        # loop over the face detections
+        for (j, rect) in enumerate(rects):
+            """
+            Determine the facial landmarks for the face region, then convert the facial landmark (x, y)-coordinates to a NumPy array
+            """
+            shape = self.predictor(im, rect)
+            shape = face_utils.shape_to_np(shape)
+
+            # we extract the face
+            vertices = cv2.convexHull(shape)
+            mask = np.zeros(im.shape[:2],np.uint8)
+            cv2.fillConvexPoly(mask, vertices, 1)
+            if self.use_grabcut:
+                bgdModel = np.zeros((1,65),np.float64)
+                fgdModel = np.zeros((1,65),np.float64)
+                rect = (0,0,im.shape[1],im.shape[2])
+                (x,y),radius = cv2.minEnclosingCircle(vertices)
+                center = (int(x),int(y))
+                radius = int(radius*self.scale_mask)
+                mask = cv2.circle(mask,center,radius,cv2.GC_PR_FGD,-1)
+                cv2.fillConvexPoly(mask, vertices, cv2.GC_FGD)
+                cv2.grabCut(im,mask,rect,bgdModel,fgdModel,5,cv2.GC_INIT_WITH_MASK)
+                mask = np.where((mask==2)|(mask==0),0,1)
+            return mask
+
+    def set_reference_images(self, images_list):
+        assert(len(images_list) != 0 and len(images_list) <= self.batch_size)
+        loaded_image = load_images(images_list, self.img_size, sharpen=self.sharpen_input)
+        image_features = None
+        if self.perceptual_model is not None:
+            image_features = self.perceptual_model.predict_on_batch(preprocess_input(np.array(loaded_image)))
+            weight_mask = np.ones(self.features_weight.shape)
+
+        if self.face_mask:
+            image_mask = np.zeros(self.ref_weight.shape)
+            for (i, im) in enumerate(loaded_image):
+                try:
+                    _, img_name = os.path.split(images_list[i])
+                    mask_img = os.path.join(self.mask_dir, f'{img_name}')
+                    if (os.path.isfile(mask_img)):
+                        print("Loading mask " + mask_img)
+                        imask = PIL.Image.open(mask_img).convert('L')
+                        mask = np.array(imask)/255
+                        mask = np.expand_dims(mask,axis=-1)
+                    else:
+                        mask = self.generate_face_mask(im)
+                        imask = (255*mask).astype('uint8')
+                        imask = PIL.Image.fromarray(imask, 'L')
+                        print("Saving mask " + mask_img)
+                        imask.save(mask_img, 'PNG')
+                        mask = np.expand_dims(mask,axis=-1)
+                    mask = np.ones(im.shape,np.float32) * mask
+                except Exception as e:
+                    print("Exception in mask handling for " + mask_img)
+                    traceback.print_exc()
+                    mask = np.ones(im.shape[:2],np.uint8)
+                    mask = np.ones(im.shape,np.float32) * np.expand_dims(mask,axis=-1)
+                image_mask[i] = mask
+            img = None
+        else:
+            image_mask = np.ones(self.ref_weight.shape)
+
+        if len(images_list) != self.batch_size:
+            if image_features is not None:
+                features_space = list(self.features_weight.shape[1:])
+                existing_features_shape = [len(images_list)] + features_space
+                empty_features_shape = [self.batch_size - len(images_list)] + features_space
+                existing_examples = np.ones(shape=existing_features_shape)
+                empty_examples = np.zeros(shape=empty_features_shape)
+                weight_mask = np.vstack([existing_examples, empty_examples])
+                image_features = np.vstack([image_features, np.zeros(empty_features_shape)])
+
+            images_space = list(self.ref_weight.shape[1:])
+            existing_images_space = [len(images_list)] + images_space
+            empty_images_space = [self.batch_size - len(images_list)] + images_space
+            existing_images = np.ones(shape=existing_images_space)
+            empty_images = np.zeros(shape=empty_images_space)
+            image_mask = image_mask * np.vstack([existing_images, empty_images])
+            loaded_image = np.vstack([loaded_image, np.zeros(empty_images_space)])
+
+        if image_features is not None:
+            self.assign_placeholder("features_weight", weight_mask)
+            self.assign_placeholder("ref_img_features", image_features)
+        self.assign_placeholder("ref_weight", image_mask)
+        self.assign_placeholder("ref_img", loaded_image)
+
+    def optimize(self, vars_to_optimize, iterations=200, use_optimizer='adam'):
+        vars_to_optimize = vars_to_optimize if isinstance(vars_to_optimize, list) else [vars_to_optimize]
+        if use_optimizer == 'lbfgs':
+            optimizer = tf.contrib.opt.ScipyOptimizerInterface(self.loss, var_list=vars_to_optimize, method='L-BFGS-B', options={'maxiter': iterations})
+        else:
+            if use_optimizer == 'ggt':
+                optimizer = tf.contrib.opt.GGTOptimizer(learning_rate=self.learning_rate)
+            else:
+                optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
+            min_op = optimizer.minimize(self.loss, var_list=[vars_to_optimize])
+            self.sess.run(tf.variables_initializer(optimizer.variables()))
+            fetch_ops = [min_op, self.loss, self.learning_rate]
+        #min_op = optimizer.minimize(self.sess)
+        #optim_results = tfp.optimizer.lbfgs_minimize(make_val_and_grad_fn(get_loss), initial_position=vars_to_optimize, num_correction_pairs=10, tolerance=1e-8)
+        self.sess.run(self._reset_global_step)
+        #self.sess.graph.finalize()  # Graph is read-only after this statement.
+        for _ in range(iterations):
+            if use_optimizer == 'lbfgs':
+                optimizer.minimize(self.sess, fetches=[vars_to_optimize, self.loss])
+                yield {"loss":self.loss.eval()}
+            else:
+                _, loss, lr = self.sess.run(fetch_ops)
+                yield {"loss":loss,"lr":lr}