Delete dnnlib/tflib/network.py

dnnlib/tflib/network.py

@@ -1,781 +0,0 @@
-# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
-#
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
-
-"""Helper for managing networks."""
-
-import types
-import inspect
-import re
-import uuid
-import sys
-import copy
-import numpy as np
-import tensorflow as tf
-
-from collections import OrderedDict
-from typing import Any, List, Tuple, Union, Callable
-
-from . import tfutil
-from .. import util
-
-from .tfutil import TfExpression, TfExpressionEx
-
-# pylint: disable=protected-access
-# pylint: disable=attribute-defined-outside-init
-# pylint: disable=too-many-public-methods
-
-_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. Defaults to build func name if None.
-        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.
-    """
-
-    def __init__(self, name: str = None, func_name: Any = None, **static_kwargs):
-        # Locate the user-specified build function.
-        assert isinstance(func_name, str) or util.is_top_level_function(func_name)
-        if util.is_top_level_function(func_name):
-            func_name = util.get_top_level_function_name(func_name)
-        module, func_name = util.get_module_from_obj_name(func_name)
-        func = util.get_obj_from_module(module, func_name)
-
-        # Dig up source code for the module containing the build function.
-        module_src = _import_module_src.get(module, None)
-        if module_src is None:
-            module_src = inspect.getsource(module)
-
-        # Initialize fields.
-        self._init_fields(name=(name or func_name), static_kwargs=static_kwargs, build_func=func, build_func_name=func_name, build_module_src=module_src)
-
-    def _init_fields(self, name: str, static_kwargs: dict, build_func: Callable, build_func_name: str, build_module_src: str) -> None:
-        tfutil.assert_tf_initialized()
-        assert isinstance(name, str)
-        assert len(name) >= 1
-        assert re.fullmatch(r"[A-Za-z0-9_.\\-]*", name)
-        assert isinstance(static_kwargs, dict)
-        assert util.is_pickleable(static_kwargs)
-        assert callable(build_func)
-        assert isinstance(build_func_name, str)
-        assert isinstance(build_module_src, str)
-
-        # Choose TensorFlow name scope.
-        with tf.name_scope(None):
-            scope = tf.get_default_graph().unique_name(name, mark_as_used=True)
-
-        # Query current TensorFlow device.
-        with tfutil.absolute_name_scope(scope), tf.control_dependencies(None):
-            device = tf.no_op(name="_QueryDevice").device
-
-        # Immutable state.
-        self._name                  = name
-        self._scope                 = scope
-        self._device                = device
-        self._static_kwargs         = util.EasyDict(copy.deepcopy(static_kwargs))
-        self._build_func            = build_func
-        self._build_func_name       = build_func_name
-        self._build_module_src      = build_module_src
-
-        # State before _init_graph().
-        self._var_inits             = dict()    # var_name => initial_value, set to None by _init_graph()
-        self._all_inits_known       = False     # Do we know for sure that _var_inits covers all the variables?
-        self._components            = None      # subnet_name => Network, None if the components are not known yet
-
-        # Initialized by _init_graph().
-        self._input_templates       = None
-        self._output_templates      = None
-        self._own_vars              = None
-
-        # Cached values initialized the respective methods.
-        self._input_shapes          = None
-        self._output_shapes         = None
-        self._input_names           = None
-        self._output_names          = None
-        self._vars                  = None
-        self._trainables            = None
-        self._var_global_to_local   = None
-        self._run_cache             = dict()
-
-    def _init_graph(self) -> None:
-        assert self._var_inits is not None
-        assert self._input_templates is None
-        assert self._output_templates is None
-        assert self._own_vars is None
-
-        # Initialize components.
-        if self._components is None:
-            self._components = util.EasyDict()
-
-        # Choose build func kwargs.
-        build_kwargs = dict(self.static_kwargs)
-        build_kwargs["is_template_graph"] = True
-        build_kwargs["components"] = self._components
-
-        # Override scope and device, and ignore surrounding control dependencies.
-        with tfutil.absolute_variable_scope(self.scope, reuse=False), tfutil.absolute_name_scope(self.scope), tf.device(self.device), tf.control_dependencies(None):
-            assert tf.get_variable_scope().name == self.scope
-            assert tf.get_default_graph().get_name_scope() == self.scope
-
-            # Create input templates.
-            self._input_templates = []
-            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_templates.append(tf.placeholder(tf.float32, name=param.name))
-
-            # Call build func.
-            out_expr = self._build_func(*self._input_templates, **build_kwargs)
-
-        # Collect output templates and variables.
-        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._own_vars = OrderedDict((var.name[len(self.scope) + 1:].split(":")[0], var) for var in tf.global_variables(self.scope + "/"))
-
-        # Check for errors.
-        if len(self._input_templates) == 0:
-            raise ValueError("Network build func did not list any inputs.")
-        if len(self._output_templates) == 0:
-            raise ValueError("Network build func did not return any outputs.")
-        if any(not tfutil.is_tf_expression(t) for t in self._output_templates):
-            raise ValueError("Network outputs must be TensorFlow expressions.")
-        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.")
-
-        # Initialize variables.
-        if len(self._var_inits):
-            tfutil.set_vars({self._get_vars()[name]: value for name, value in self._var_inits.items() if name in self._get_vars()})
-        remaining_inits = [var.initializer for name, var in self._own_vars.items() if name not in self._var_inits]
-        if self._all_inits_known:
-            assert len(remaining_inits) == 0
-        else:
-            tfutil.run(remaining_inits)
-        self._var_inits = None
-
-    @property
-    def name(self):
-        """User-specified name string."""
-        return self._name
-
-    @property
-    def scope(self):
-        """Unique TensorFlow scope containing template graph and variables, derived from the user-specified name."""
-        return self._scope
-
-    @property
-    def device(self):
-        """Name of the TensorFlow device that the weights of this network reside on. Determined by the current device at construction time."""
-        return self._device
-
-    @property
-    def static_kwargs(self):
-        """EasyDict of arguments passed to the user-supplied build func."""
-        return copy.deepcopy(self._static_kwargs)
-
-    @property
-    def components(self):
-        """EasyDict of sub-networks created by the build func."""
-        return copy.copy(self._get_components())
-
-    def _get_components(self):
-        if self._components is None:
-            self._init_graph()
-            assert self._components is not None
-        return self._components
-
-    @property
-    def input_shapes(self):
-        """List of input tensor shapes, including minibatch dimension."""
-        if self._input_shapes is None:
-            self._input_shapes = [t.shape.as_list() for t in self.input_templates]
-        return copy.deepcopy(self._input_shapes)
-
-    @property
-    def output_shapes(self):
-        """List of output tensor shapes, including minibatch dimension."""
-        if self._output_shapes is None:
-            self._output_shapes = [t.shape.as_list() for t in self.output_templates]
-        return copy.deepcopy(self._output_shapes)
-
-    @property
-    def input_shape(self):
-        """Short-hand for input_shapes[0]."""
-        return self.input_shapes[0]
-
-    @property
-    def output_shape(self):
-        """Short-hand for output_shapes[0]."""
-        return self.output_shapes[0]
-
-    @property
-    def num_inputs(self):
-        """Number of input tensors."""
-        return len(self.input_shapes)
-
-    @property
-    def num_outputs(self):
-        """Number of output tensors."""
-        return len(self.output_shapes)
-
-    @property
-    def input_names(self):
-        """Name string for each input."""
-        if self._input_names is None:
-            self._input_names = [t.name.split("/")[-1].split(":")[0] for t in self.input_templates]
-        return copy.copy(self._input_names)
-
-    @property
-    def output_names(self):
-        """Name string for each output."""
-        if self._output_names is None:
-            self._output_names = [t.name.split("/")[-1].split(":")[0] for t in self.output_templates]
-        return copy.copy(self._output_names)
-
-    @property
-    def input_templates(self):
-        """Input placeholders in the template graph."""
-        if self._input_templates is None:
-            self._init_graph()
-            assert self._input_templates is not None
-        return copy.copy(self._input_templates)
-
-    @property
-    def output_templates(self):
-        """Output tensors in the template graph."""
-        if self._output_templates is None:
-            self._init_graph()
-            assert self._output_templates is not None
-        return copy.copy(self._output_templates)
-
-    @property
-    def own_vars(self):
-        """Variables defined by this network (local_name => var), excluding sub-networks."""
-        return copy.copy(self._get_own_vars())
-
-    def _get_own_vars(self):
-        if self._own_vars is None:
-            self._init_graph()
-            assert self._own_vars is not None
-        return self._own_vars
-
-    @property
-    def vars(self):
-        """All variables (local_name => var)."""
-        return copy.copy(self._get_vars())
-
-    def _get_vars(self):
-        if self._vars is None:
-            self._vars = OrderedDict(self._get_own_vars())
-            for comp in self._get_components().values():
-                self._vars.update((comp.name + "/" + name, var) for name, var in comp._get_vars().items())
-        return self._vars
-
-    @property
-    def trainables(self):
-        """All trainable variables (local_name => var)."""
-        return copy.copy(self._get_trainables())
-
-    def _get_trainables(self):
-        if self._trainables is None:
-            self._trainables = OrderedDict((name, var) for name, var in self.vars.items() if var.trainable)
-        return self._trainables
-
-    @property
-    def var_global_to_local(self):
-        """Mapping from variable global names to local names."""
-        return copy.copy(self._get_var_global_to_local())
-
-    def _get_var_global_to_local(self):
-        if self._var_global_to_local is None:
-            self._var_global_to_local = OrderedDict((var.name.split(":")[0], name) for name, var in self.vars.items())
-        return self._var_global_to_local
-
-    def reset_own_vars(self) -> None:
-        """Re-initialize all variables of this network, excluding sub-networks."""
-        if self._var_inits is None or self._components is None:
-            tfutil.run([var.initializer for var in self._get_own_vars().values()])
-        else:
-            self._var_inits.clear()
-            self._all_inits_known = False
-
-    def reset_vars(self) -> None:
-        """Re-initialize all variables of this network, including sub-networks."""
-        if self._var_inits is None:
-            tfutil.run([var.initializer for var in self._get_vars().values()])
-        else:
-            self._var_inits.clear()
-            self._all_inits_known = False
-            if self._components is not None:
-                for comp in self._components.values():
-                    comp.reset_vars()
-
-    def reset_trainables(self) -> None:
-        """Re-initialize all trainable variables of this network, including sub-networks."""
-        tfutil.run([var.initializer for var in self._get_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).
-        The graph is placed on the current TensorFlow device."""
-        assert len(in_expr) == self.num_inputs
-        assert not all(expr is None for expr in in_expr)
-        self._get_vars()  # ensure that all variables have been created
-
-        # Choose 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.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._get_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._get_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"]            = 5
-        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._get_own_vars().keys(), tfutil.run(list(self._get_own_vars().values()))))
-        state["input_shapes"]       = self.input_shapes
-        state["output_shapes"]      = self.output_shapes
-        state["input_names"]        = self.input_names
-        state["output_names"]       = self.output_names
-        return state
-
-    def __setstate__(self, state: dict) -> None:
-        """Pickle import."""
-
-        # Execute custom import handlers.
-        for handler in _import_handlers:
-            state = handler(state)
-
-        # Get basic fields.
-        assert state["version"] in [2, 3, 4, 5]
-        name = state["name"]
-        static_kwargs = state["static_kwargs"]
-        build_module_src = state["build_module_src"]
-        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] = build_module_src
-        exec(build_module_src, module.__dict__) # pylint: disable=exec-used
-        build_func = util.get_obj_from_module(module, build_func_name)
-
-        # Initialize fields.
-        self._init_fields(name=name, static_kwargs=static_kwargs, build_func=build_func, build_func_name=build_func_name, build_module_src=build_module_src)
-        self._var_inits.update(copy.deepcopy(state["variables"]))
-        self._all_inits_known   = True
-        self._components        = util.EasyDict(state.get("components", {}))
-        self._input_shapes      = copy.deepcopy(state.get("input_shapes", None))
-        self._output_shapes     = copy.deepcopy(state.get("output_shapes", None))
-        self._input_names       = copy.deepcopy(state.get("input_names", None))
-        self._output_names      = copy.deepcopy(state.get("output_names", None))
-
-    def clone(self, name: str = None, **new_static_kwargs) -> "Network":
-        """Create a clone of this network with its own copy of the variables."""
-        static_kwargs = dict(self.static_kwargs)
-        static_kwargs.update(new_static_kwargs)
-        net = object.__new__(Network)
-        net._init_fields(name=(name or self.name), static_kwargs=static_kwargs, build_func=self._build_func, build_func_name=self._build_func_name, build_module_src=self._build_module_src)
-        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."""
-
-        # Source has unknown variables or unknown components => init now.
-        if (src_net._var_inits is not None and not src_net._all_inits_known) or src_net._components is None:
-            src_net._get_vars()
-
-       # Both networks are inited => copy directly.
-        if src_net._var_inits is None and self._var_inits is None:
-            names = [name for name in self._get_own_vars().keys() if name in src_net._get_own_vars()]
-            tfutil.set_vars(tfutil.run({self._get_vars()[name]: src_net._get_vars()[name] for name in names}))
-            return
-
-        # Read from source.
-        if src_net._var_inits is None:
-            value_dict = tfutil.run(src_net._get_own_vars())
-        else:
-            value_dict = src_net._var_inits
-
-        # Write to destination.
-        if self._var_inits is None:
-            tfutil.set_vars({self._get_vars()[name]: value for name, value in value_dict.items() if name in self._get_vars()})
-        else:
-            self._var_inits.update(value_dict)
-
-    def copy_vars_from(self, src_net: "Network") -> None:
-        """Copy the values of all variables from the given network, including sub-networks."""
-
-        # Source has unknown variables or unknown components => init now.
-        if (src_net._var_inits is not None and not src_net._all_inits_known) or src_net._components is None:
-            src_net._get_vars()
-
-        # Source is inited, but destination components have not been created yet => set as initial values.
-        if src_net._var_inits is None and self._components is None:
-            self._var_inits.update(tfutil.run(src_net._get_vars()))
-            return
-
-        # Destination has unknown components => init now.
-        if self._components is None:
-            self._get_vars()
-
-        # Both networks are inited => copy directly.
-        if src_net._var_inits is None and self._var_inits is None:
-            names = [name for name in self._get_vars().keys() if name in src_net._get_vars()]
-            tfutil.set_vars(tfutil.run({self._get_vars()[name]: src_net._get_vars()[name] for name in names}))
-            return
-
-        # Copy recursively, component by component.
-        self.copy_own_vars_from(src_net)
-        for name, src_comp in src_net._components.items():
-            if name in self._components:
-                self._components[name].copy_vars_from(src_comp)
-
-    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._get_trainables().keys() if name in src_net._get_trainables()]
-        tfutil.set_vars(tfutil.run({self._get_vars()[name]: src_net._get_vars()[name] 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._get_vars().items():
-                if name in src_net._get_vars():
-                    cur_beta = beta if var.trainable else beta_nontrainable
-                    new_value = tfutil.lerp(src_net._get_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,
-            **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.
-        """
-        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):
-                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(self.device if num_gpus == 1 else "/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] + expr.shape.as_list()[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.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]:
-        _ = self.output_templates  # ensure that the template graph has been created
-        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):
-            if len(parent_ops) == 0 and len(parent_vars) == 0:
-                return
-
-            # 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 not in ["Identity", "Cast", "Transpose"] for op in cur_ops)
-            if (level == 0 or not contains_direct_ops) and (len(cur_ops) != 0 or len(cur_vars) != 0):
-                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._get_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(int(np.prod(var.shape.as_list())) for var in layer_trainables)
-            weights = [var for var in layer_trainables if var.name.endswith("/weight: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._get_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