(self, loss: TfExpression, trainable_vars: Union[List, dict])
Register the gradients of the given loss function with respect to the given variables.
        Intended to be called once per GPU.
Register the gradients of the given loss function with respect to the given variables.
        Intended to be called once per GPU.
Register the gradients of the given loss function with respect to the given variables . Intended to be called once per GPU .
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 = []
        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 register_gradients ( self , loss : TfExpression , trainable_vars : Union [ List , dict ] ) -> None : 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 = [ ] 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 )