Spaces:
Runtime error
Runtime error
| # 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) | |