Delete dnnlib/tflib/optimizer.py

dnnlib/tflib/optimizer.py

@@ -1,372 +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 wrapper for a Tensorflow optimizer."""
-
-import platform
-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
-
-_collective_ops_warning_printed = False
-_collective_ops_group_key       = 831766147
-_collective_ops_instance_key    = 436340067
-
-class Optimizer:
-    """A Wrapper for tf.train.Optimizer.
-
-    Automatically takes care of:
-    - Gradient averaging for multi-GPU training.
-    - Gradient accumulation for arbitrarily large minibatches.
-    - 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",                  # Name string that will appear in TensorFlow graph.
-        tf_optimizer:           str             = "tf.train.AdamOptimizer", # Underlying optimizer class.
-        learning_rate:          TfExpressionEx  = 0.001,                    # Learning rate. Can vary over time.
-        minibatch_multiplier:   TfExpressionEx  = None,                     # Treat N consecutive minibatches as one by accumulating gradients.
-        share:                  "Optimizer"     = None,                     # Share internal state with a previously created optimizer?
-        use_loss_scaling:       bool            = False,                    # Enable dynamic loss scaling for robust mixed-precision training?
-        loss_scaling_init:      float           = 64.0,                     # Log2 of initial loss scaling factor.
-        loss_scaling_inc:       float           = 0.0005,                   # Log2 of per-minibatch loss scaling increment when there is no overflow.
-        loss_scaling_dec:       float           = 1.0,                      # Log2 of per-minibatch loss scaling decrement when there is an overflow.
-        report_mem_usage:       bool            = False,                    # Report fine-grained memory usage statistics in TensorBoard?
-        **kwargs):
-
-        # Public fields.
-        self.name                   = name
-        self.learning_rate          = learning_rate
-        self.minibatch_multiplier   = minibatch_multiplier
-        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
-
-        # Private fields.
-        self._updates_applied       = False
-        self._devices               = OrderedDict() # device_name => EasyDict()
-        self._shared_optimizers     = OrderedDict() # device_name => optimizer_class
-        self._gradient_shapes       = None          # [shape, ...]
-        self._report_mem_usage      = report_mem_usage
-
-        # Validate arguments.
-        assert callable(self.optimizer_class)
-
-        # Share internal state if requested.
-        if share is not None:
-            assert isinstance(share, Optimizer)
-            assert self.optimizer_class is share.optimizer_class
-            assert self.learning_rate is share.learning_rate
-            assert self.optimizer_kwargs == share.optimizer_kwargs
-            self._shared_optimizers = share._shared_optimizers # pylint: disable=protected-access
-
-    def _get_device(self, device_name: str):
-        """Get internal state for the given TensorFlow device."""
-        tfutil.assert_tf_initialized()
-        if device_name in self._devices:
-            return self._devices[device_name]
-
-        # Initialize fields.
-        device = util.EasyDict()
-        device.name             = device_name
-        device.optimizer        = None          # Underlying optimizer:     optimizer_class
-        device.loss_scaling_var = None          # Log2 of loss scaling:     tf.Variable
-        device.grad_raw         = OrderedDict() # Raw gradients:            var => [grad, ...]
-        device.grad_clean       = OrderedDict() # Clean gradients:          var => grad
-        device.grad_acc_vars    = OrderedDict() # Accumulation sums:        var => tf.Variable
-        device.grad_acc_count   = None          # Accumulation counter:     tf.Variable
-        device.grad_acc         = OrderedDict() # Accumulated gradients:    var => grad
-
-        # Setup TensorFlow objects.
-        with tfutil.absolute_name_scope(self.scope + "/Devices"), tf.device(device_name), tf.control_dependencies(None):
-            if device_name not in self._shared_optimizers:
-                optimizer_name = self.scope.replace("/", "_") + "_opt%d" % len(self._shared_optimizers)
-                self._shared_optimizers[device_name] = self.optimizer_class(name=optimizer_name, learning_rate=self.learning_rate, **self.optimizer_kwargs)
-            device.optimizer = self._shared_optimizers[device_name]
-            if self.use_loss_scaling:
-                device.loss_scaling_var = tf.Variable(np.float32(self.loss_scaling_init), trainable=False, name="loss_scaling_var")
-
-        # Register device.
-        self._devices[device_name] = device
-        return device
-
-    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."""
-        tfutil.assert_tf_initialized()
-        assert not self._updates_applied
-        device = self._get_device(loss.device)
-
-        # Validate trainables.
-        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])
-        assert all(var.device == device.name for var in trainable_vars)
-
-        # Validate shapes.
-        if self._gradient_shapes is None:
-            self._gradient_shapes = [var.shape.as_list() for var in trainable_vars]
-        assert len(trainable_vars) == len(self._gradient_shapes)
-        assert all(var.shape.as_list() == var_shape for var, var_shape in zip(trainable_vars, self._gradient_shapes))
-
-        # Report memory usage if requested.
-        deps = [loss]
-        if self._report_mem_usage:
-            self._report_mem_usage = False
-            try:
-                with tf.name_scope(self.id + '_mem'), tf.device(device.name), tf.control_dependencies([loss]):
-                    deps.append(autosummary.autosummary(self.id + "/mem_usage_gb", tf.contrib.memory_stats.BytesInUse() / 2**30))
-            except tf.errors.NotFoundError:
-                pass
-
-        # Compute gradients.
-        with tf.name_scope(self.id + "_grad"), tf.device(device.name), tf.control_dependencies(deps):
-            loss = self.apply_loss_scaling(tf.cast(loss, tf.float32))
-            gate = tf.train.Optimizer.GATE_NONE  # disable gating to reduce memory usage
-            grad_list = device.optimizer.compute_gradients(loss=loss, var_list=trainable_vars, gate_gradients=gate)
-
-        # Register gradients.
-        for grad, var in grad_list:
-            if var not in device.grad_raw:
-                device.grad_raw[var] = []
-            device.grad_raw[var].append(grad)
-
-    def apply_updates(self, allow_no_op: bool = False) -> 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
-        all_ops = []
-
-        # Check for no-op.
-        if allow_no_op and len(self._devices) == 0:
-            with tfutil.absolute_name_scope(self.scope):
-                return tf.no_op(name='TrainingOp')
-
-        # Clean up gradients.
-        for device_idx, device in enumerate(self._devices.values()):
-            with tfutil.absolute_name_scope(self.scope + "/Clean%d" % device_idx), tf.device(device.name):
-                for var, grad in device.grad_raw.items():
-
-                    # Filter out disconnected gradients and convert to float32.
-                    grad = [g for g in grad if g is not None]
-                    grad = [tf.cast(g, tf.float32) for g in grad]
-
-                    # Sum within the device.
-                    if len(grad) == 0:
-                        grad = tf.zeros(var.shape)  # No gradients => zero.
-                    elif len(grad) == 1:
-                        grad = grad[0]              # Single gradient => use as is.
-                    else:
-                        grad = tf.add_n(grad)       # Multiple gradients => sum.
-
-                    # Scale as needed.
-                    scale = 1.0 / len(device.grad_raw[var]) / len(self._devices)
-                    scale = tf.constant(scale, dtype=tf.float32, name="scale")
-                    if self.minibatch_multiplier is not None:
-                        scale /= tf.cast(self.minibatch_multiplier, tf.float32)
-                    scale = self.undo_loss_scaling(scale)
-                    device.grad_clean[var] = grad * scale
-
-        # Sum gradients across devices.
-        if len(self._devices) > 1:
-            with tfutil.absolute_name_scope(self.scope + "/Broadcast"), tf.device(None):
-                if platform.system() == "Windows":    # Windows => NCCL ops are not available.
-                    self._broadcast_fallback()
-                elif tf.VERSION.startswith("1.15."):  # TF 1.15 => NCCL ops are broken: https://github.com/tensorflow/tensorflow/issues/41539
-                    self._broadcast_fallback()
-                else:                                 # Otherwise => NCCL ops are safe to use.
-                    self._broadcast_nccl()
-
-        # Apply updates separately on each device.
-        for device_idx, device in enumerate(self._devices.values()):
-            with tfutil.absolute_name_scope(self.scope + "/Apply%d" % device_idx), tf.device(device.name):
-                # pylint: disable=cell-var-from-loop
-
-                # Accumulate gradients over time.
-                if self.minibatch_multiplier is None:
-                    acc_ok = tf.constant(True, name='acc_ok')
-                    device.grad_acc = OrderedDict(device.grad_clean)
-                else:
-                    # Create variables.
-                    with tf.control_dependencies(None):
-                        for var in device.grad_clean.keys():
-                            device.grad_acc_vars[var] = tf.Variable(tf.zeros(var.shape), trainable=False, name="grad_acc_var")
-                        device.grad_acc_count = tf.Variable(tf.zeros([]), trainable=False, name="grad_acc_count")
-
-                    # Track counter.
-                    count_cur = device.grad_acc_count + 1.0
-                    count_inc_op = lambda: tf.assign(device.grad_acc_count, count_cur)
-                    count_reset_op = lambda: tf.assign(device.grad_acc_count, tf.zeros([]))
-                    acc_ok = (count_cur >= tf.cast(self.minibatch_multiplier, tf.float32))
-                    all_ops.append(tf.cond(acc_ok, count_reset_op, count_inc_op))
-
-                    # Track gradients.
-                    for var, grad in device.grad_clean.items():
-                        acc_var = device.grad_acc_vars[var]
-                        acc_cur = acc_var + grad
-                        device.grad_acc[var] = acc_cur
-                        with tf.control_dependencies([acc_cur]):
-                            acc_inc_op = lambda: tf.assign(acc_var, acc_cur)
-                            acc_reset_op = lambda: tf.assign(acc_var, tf.zeros(var.shape))
-                            all_ops.append(tf.cond(acc_ok, acc_reset_op, acc_inc_op))
-
-                # No overflow => apply gradients.
-                all_ok = tf.reduce_all(tf.stack([acc_ok] + [tf.reduce_all(tf.is_finite(g)) for g in device.grad_acc.values()]))
-                apply_op = lambda: device.optimizer.apply_gradients([(tf.cast(grad, var.dtype), var) for var, grad in device.grad_acc.items()])
-                all_ops.append(tf.cond(all_ok, apply_op, tf.no_op))
-
-                # Adjust loss scaling.
-                if self.use_loss_scaling:
-                    ls_inc_op = lambda: tf.assign_add(device.loss_scaling_var, self.loss_scaling_inc)
-                    ls_dec_op = lambda: tf.assign_sub(device.loss_scaling_var, self.loss_scaling_dec)
-                    ls_update_op = lambda: tf.group(tf.cond(all_ok, ls_inc_op, ls_dec_op))
-                    all_ops.append(tf.cond(acc_ok, ls_update_op, tf.no_op))
-
-                # Last device => report statistics.
-                if device_idx == len(self._devices) - 1:
-                    all_ops.append(autosummary.autosummary(self.id + "/learning_rate", tf.convert_to_tensor(self.learning_rate)))
-                    all_ops.append(autosummary.autosummary(self.id + "/overflow_frequency", tf.where(all_ok, 0, 1), condition=acc_ok))
-                    if self.use_loss_scaling:
-                        all_ops.append(autosummary.autosummary(self.id + "/loss_scaling_log2", device.loss_scaling_var))
-
-        # Initialize variables.
-        self.reset_optimizer_state()
-        if self.use_loss_scaling:
-            tfutil.init_uninitialized_vars([device.loss_scaling_var for device in self._devices.values()])
-        if self.minibatch_multiplier is not None:
-            tfutil.run([var.initializer for device in self._devices.values() for var in list(device.grad_acc_vars.values()) + [device.grad_acc_count]])
-
-        # Group everything into a single op.
-        with tfutil.absolute_name_scope(self.scope):
-            return tf.group(*all_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 device in self._devices.values() for var in device.optimizer.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."""
-        return self._get_device(device).loss_scaling_var
-
-    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
-
-    def _broadcast_nccl(self):
-        """Sum gradients across devices using NCCL ops (fast path)."""
-        from tensorflow.python.ops import nccl_ops # pylint: disable=no-name-in-module
-        for all_vars in zip(*[device.grad_clean.keys() for device in self._devices.values()]):
-            if any(x.shape.num_elements() > 0 for x in all_vars):
-                all_grads = [device.grad_clean[var] for device, var in zip(self._devices.values(), all_vars)]
-                all_grads = nccl_ops.all_sum(all_grads)
-                for device, var, grad in zip(self._devices.values(), all_vars, all_grads):
-                    device.grad_clean[var] = grad
-
-    def _broadcast_fallback(self):
-        """Sum gradients across devices using TensorFlow collective ops (slow fallback path)."""
-        from tensorflow.python.ops import collective_ops # pylint: disable=no-name-in-module
-        global _collective_ops_warning_printed, _collective_ops_group_key, _collective_ops_instance_key
-        if all(x.shape.num_elements() == 0 for device in self._devices.values() for x in device.grad_clean.values()):
-            return
-        if not _collective_ops_warning_printed:
-            print("------------------------------------------------------------------------")
-            print("WARNING: Using slow fallback implementation for inter-GPU communication.")
-            print("Please use TensorFlow 1.14 on Linux for optimal training performance.")
-            print("------------------------------------------------------------------------")
-            _collective_ops_warning_printed = True
-        for device in self._devices.values():
-            with tf.device(device.name):
-                combo = [tf.reshape(x, [x.shape.num_elements()]) for x in device.grad_clean.values()]
-                combo = tf.concat(combo, axis=0)
-                combo = collective_ops.all_reduce(combo, merge_op='Add', final_op='Id',
-                    group_size=len(self._devices), group_key=_collective_ops_group_key,
-                    instance_key=_collective_ops_instance_key)
-                cur_ofs = 0
-                for var, grad_old in device.grad_clean.items():
-                    grad_new = tf.reshape(combo[cur_ofs : cur_ofs + grad_old.shape.num_elements()], grad_old.shape)
-                    cur_ofs += grad_old.shape.num_elements()
-                    device.grad_clean[var] = grad_new
-        _collective_ops_instance_key += 1
-
-
-class SimpleAdam:
-    """Simplified version of tf.train.AdamOptimizer that behaves identically when used with dnnlib.tflib.Optimizer."""
-
-    def __init__(self, name="Adam", learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8):
-        self.name = name
-        self.learning_rate = learning_rate
-        self.beta1 = beta1
-        self.beta2 = beta2
-        self.epsilon = epsilon
-        self.all_state_vars = []
-
-    def variables(self):
-        return self.all_state_vars
-
-    def compute_gradients(self, loss, var_list, gate_gradients=tf.train.Optimizer.GATE_NONE):
-        assert gate_gradients == tf.train.Optimizer.GATE_NONE
-        return list(zip(tf.gradients(loss, var_list), var_list))
-
-    def apply_gradients(self, grads_and_vars):
-        with tf.name_scope(self.name):
-            state_vars = []
-            update_ops = []
-
-            # Adjust learning rate to deal with startup bias.
-            with tf.control_dependencies(None):
-                b1pow_var = tf.Variable(dtype=tf.float32, initial_value=1, trainable=False)
-                b2pow_var = tf.Variable(dtype=tf.float32, initial_value=1, trainable=False)
-                state_vars += [b1pow_var, b2pow_var]
-            b1pow_new = b1pow_var * self.beta1
-            b2pow_new = b2pow_var * self.beta2
-            update_ops += [tf.assign(b1pow_var, b1pow_new), tf.assign(b2pow_var, b2pow_new)]
-            lr_new = self.learning_rate * tf.sqrt(1 - b2pow_new) / (1 - b1pow_new)
-
-            # Construct ops to update each variable.
-            for grad, var in grads_and_vars:
-                with tf.control_dependencies(None):
-                    m_var = tf.Variable(dtype=tf.float32, initial_value=tf.zeros_like(var), trainable=False)
-                    v_var = tf.Variable(dtype=tf.float32, initial_value=tf.zeros_like(var), trainable=False)
-                    state_vars += [m_var, v_var]
-                m_new = self.beta1 * m_var + (1 - self.beta1) * grad
-                v_new = self.beta2 * v_var + (1 - self.beta2) * tf.square(grad)
-                var_delta = lr_new * m_new / (tf.sqrt(v_new) + self.epsilon)
-                update_ops += [tf.assign(m_var, m_new), tf.assign(v_var, v_new), tf.assign_sub(var, var_delta)]
-
-            # Group everything together.
-            self.all_state_vars += state_vars
-            return tf.group(*update_ops)