<commit_before>import numpy
import tensorflow as tf

"""Helper class for training using multiple GPUs.

Given a set of model replicas and an optimizer, it takes care of splitting
minibatches, feeding them to the individual replicas, and then combining and
applying the resulting gradients (and losses).
"""
class ModelUpdater(object):
    def __init__(self, config, num_gpus, replicas, optimizer, global_step,
                 summary_writer=None):
        assert len(replicas) > 0
        assert (len(replicas) == num_gpus
                or (len(replicas) == 1 and num_gpus == 0))

        self.replicas = replicas
        self.global_step = global_step
        self.summary_writer = summary_writer

        self.replica_weights = []
        for i in range(len(self.replicas)):
            name = 'replica_weight_{}'.format(i)
            placeholder = tf.placeholder(name=name, shape=(), dtype=tf.float32)
            self.replica_weights.append(placeholder)

        weighted_losses = []
        all_grad_vars = []
        for i in range(len(self.replicas)):
            device_type = "GPU" if num_gpus > 0 else "CPU"
            device_spec = tf.DeviceSpec(device_type=device_type, device_index=i)
            with tf.device(device_spec):
                with tf.variable_scope(tf.get_variable_scope(), reuse=(i>0)):
                    if config.loss_function == "cross-entropy":
                        loss = replicas[i].loss
                    elif config.loss_function == "per-token-cross-entropy":
                        ce_per_sent = replicas[i].loss_per_sentence
                        ce_total = tf.reduce_sum(ce_per_sent)
                        num_tokens = tf.reduce_sum(replicas[i].inputs.y_mask)
                        loss = ce_total / tf.cast(num_tokens, tf.float32)
                    else:
                        assert False
                    loss = self._regularize(loss, config.decay_c,
                                            config.map_decay_c)
                    gradients = optimizer.compute_gradients(loss)
                    all_grad_vars.append(gradients)
                    weight = self.replica_weights[i]
                    weighted_losses.append(loss*weight)

        self.loss = sum(weighted_losses) / sum(self.replica_weights)

        grad_vars = self._average_gradients(all_grad_vars, self.replica_weights)
        grads, varss = zip(*grad_vars)
        clipped_grads, global_norm = tf.clip_by_global_norm(
            grads, clip_norm=config.clip_c)
        # Might be interesting to see how the global norm changes over time,
        # attach a summary?
        grad_vars = zip(clipped_grads, varss)
        self.apply_grads = optimizer.apply_gradients(
            grad_vars, global_step=self.global_step)

        tf.summary.scalar(name='mean_cost', tensor=self.loss)
        tf.summary.scalar(name='t', tensor=self.global_step)
        self.merged = tf.summary.merge_all()

    def _regularize(self, loss, decay_c, map_decay_c):
        with tf.variable_scope("loss"):
            # Optionally, add an L2 loss term.
            if decay_c > 0.0:
                l2_loss = tf.constant(0.0, dtype=tf.float32)
                l2_loss = tf.add_n([tf.nn.l2_loss(v) for v in tf.trainable_variables()]) * tf.constant(decay_c, dtype=tf.float32)
                loss += l2_loss
            # Optionally, add an L2 loss term based on a prior model.
            if map_decay_c > 0.0:
                map_l2_loss = tf.constant(0.0, dtype=tf.float32)
                map_l2_acc = []
                for v in tf.trainable_variables():
                    prior_name = 'prior/'+v.name.split(':')[0]
                    prior_v = tf.get_variable(
                        prior_name, initializer=v.initialized_value(),<commit_msg>Disable gradient clipping if clip_c <= 0.0<commit_after>import numpy
import tensorflow as tf

"""Helper class for training using multiple GPUs.

Given a set of model replicas and an optimizer, it takes care of splitting
minibatches, feeding them to the individual replicas, and then combining and
applying the resulting gradients (and losses).
"""
class ModelUpdater(object):
    def __init__(self, config, num_gpus, replicas, optimizer, global_step,
                 summary_writer=None):
        assert len(replicas) > 0
        assert (len(replicas) == num_gpus
                or (len(replicas) == 1 and num_gpus == 0))

        self.replicas = replicas
        self.global_step = global_step
        self.summary_writer = summary_writer

        self.replica_weights = []
        for i in range(len(self.replicas)):
            name = 'replica_weight_{}'.format(i)
            placeholder = tf.placeholder(name=name, shape=(), dtype=tf.float32)
            self.replica_weights.append(placeholder)

        weighted_losses = []
        all_grad_vars = []
        for i in range(len(self.replicas)):
            device_type = "GPU" if num_gpus > 0 else "CPU"
            device_spec = tf.DeviceSpec(device_type=device_type, device_index=i)
            with tf.device(device_spec):
                with tf.variable_scope(tf.get_variable_scope(), reuse=(i>0)):
                    if config.loss_function == "cross-entropy":
                        loss = replicas[i].loss
                    elif config.loss_function == "per-token-cross-entropy":
                        ce_per_sent = replicas[i].loss_per_sentence
                        ce_total = tf.reduce_sum(ce_per_sent)
                        num_tokens = tf.reduce_sum(replicas[i].inputs.y_mask)
                        loss = ce_total / tf.cast(num_tokens, tf.float32)
                    else:
                        assert False
                    loss = self._regularize(loss, config.decay_c,
                                            config.map_decay_c)
                    gradients = optimizer.compute_gradients(loss)
                    all_grad_vars.append(gradients)
                    weight = self.replica_weights[i]
                    weighted_losses.append(loss*weight)

        self.loss = sum(weighted_losses) / sum(self.replica_weights)

        grad_vars = self._average_gradients(all_grad_vars, self.replica_weights)
        if config.clip_c > 0.0:
            grads, varss = zip(*grad_vars)
            clipped_grads, global_norm = tf.clip_by_global_norm(
                grads, clip_norm=config.clip_c)
            # Might be interesting to see how the global norm changes over time,
            # attach a summary?
            grad_vars = zip(clipped_grads, varss)
        self.apply_grads = optimizer.apply_gradients(
            grad_vars, global_step=self.global_step)

        tf.summary.scalar(name='mean_cost', tensor=self.loss)
        tf.summary.scalar(name='t', tensor=self.global_step)
        self.merged = tf.summary.merge_all()

    def _regularize(self, loss, decay_c, map_decay_c):
        with tf.variable_scope("loss"):
            # Optionally, add an L2 loss term.
            if decay_c > 0.0:
                l2_loss = tf.constant(0.0, dtype=tf.float32)
                l2_loss = tf.add_n([tf.nn.l2_loss(v) for v in tf.trainable_variables()]) * tf.constant(decay_c, dtype=tf.float32)
                loss += l2_loss
            # Optionally, add an L2 loss term based on a prior model.
            if map_decay_c > 0.0:
                map_l2_loss = tf.constant(0.0, dtype=tf.float32)
                map_l2_acc = []
                for v in tf.trainable_variables():
                    prior_name = 'prior/'+v.name.split(':')[0]
                    prior_v = tf.get_variable(
                        prior_name, initializer=v.initialized_value(),
{'fraction_non_alphanumeric': 0.09390962671905698, 'fraction_numerical': 0.008906352324819908, 'mean_word_length': 1.971206225680934, 'pattern_counts': {'":': 4, '<': 4, '<?xml version=': 0, '>': 14, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2609817', 'n_tokens_mistral': 2326, 'n_tokens_neox': 2220, 'n_words': 477}