Module neuralops.trainer.methods.rdrop_r3f

R‑Drop and R3F (regularised dropout / robust representation learning).

Functions

def apply_r3f(model, config) ‑> None

Expand source code

@register_method("r3f")
def apply_r3f(model, config) -> None:
    """Apply r3f."""
    pass

Apply r3f.

def apply_rdrop(model, config) ‑> None

Expand source code

@register_method("rdrop")
def apply_rdrop(model, config) -> None:
    """Apply rdrop."""
    pass

Apply rdrop.

def r3f_forward(model, input_ids, attention_mask, token_type_ids, labels, epsilon, is_regression) ‑> Any

Expand source code

def r3f_forward(
    model, input_ids, attention_mask, token_type_ids, labels, epsilon, is_regression
) -> Any:
    """R3f forward."""
    base_kwargs = {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels,
    }
    if token_type_ids is not None:
        base_kwargs["token_type_ids"] = token_type_ids
    out_clean = model(**base_kwargs)
    logits_clean = out_clean.logits

    embed = model.get_input_embeddings()
    inputs_embeds = embed(input_ids)
    noise = torch.randn_like(inputs_embeds) * epsilon
    noise_kwargs = {
        "inputs_embeds": inputs_embeds + noise,
        "attention_mask": attention_mask,
        "labels": labels,
    }
    if token_type_ids is not None:
        noise_kwargs["token_type_ids"] = token_type_ids
    out_noisy = model(**noise_kwargs)
    logits_noisy = out_noisy.logits

    if is_regression:
        mse = nn.MSELoss()
        loss_clean = mse(logits_clean.view(-1), labels.view(-1))
        loss_noisy = mse(logits_noisy.view(-1), labels.view(-1))
        p = torch.distributions.Normal(logits_clean, torch.ones_like(logits_clean))
        q = torch.distributions.Normal(logits_noisy, torch.ones_like(logits_noisy))
        kl = torch.distributions.kl_divergence(p, q).mean()
        loss = (loss_clean + loss_noisy) / 2.0 + kl
    else:
        ce = nn.CrossEntropyLoss()
        loss_clean = ce(logits_clean.view(-1, logits_clean.size(-1)), labels.view(-1))
        loss_noisy = ce(logits_noisy.view(-1, logits_noisy.size(-1)), labels.view(-1))
        logp_clean = F.log_softmax(logits_clean, dim=-1)
        logp_noisy = F.log_softmax(logits_noisy, dim=-1)
        p = torch.softmax(logits_clean, dim=-1)
        q = torch.softmax(logits_noisy, dim=-1)
        kl1 = F.kl_div(logp_clean, q, reduction="batchmean")
        kl2 = F.kl_div(logp_noisy, p, reduction="batchmean")
        loss = (loss_clean + loss_noisy) / 2.0 + (kl1 + kl2) / 2.0
    return loss, logits_clean

R3f forward.

def rdrop_loss(logits1, logits2, labels, is_regression, alpha) ‑> Any

Expand source code

def rdrop_loss(logits1, logits2, labels, is_regression, alpha) -> Any:
    """Rdrop loss."""
    if is_regression:
        loss_fn = nn.MSELoss()
        loss1 = loss_fn(logits1.view(-1), labels.view(-1))
        loss2 = loss_fn(logits2.view(-1), labels.view(-1))
        p = torch.distributions.Normal(logits1, torch.ones_like(logits1))
        q = torch.distributions.Normal(logits2, torch.ones_like(logits2))
        kl = (
            torch.distributions.kl_divergence(p, q).mean()
            + torch.distributions.kl_divergence(q, p).mean()
        ) / 2.0
        loss = (loss1 + loss2) / 2.0 + alpha * kl
    else:
        ce = nn.CrossEntropyLoss()
        loss1 = ce(logits1.view(-1, logits1.size(-1)), labels.view(-1))
        loss2 = ce(logits2.view(-1, logits2.size(-1)), labels.view(-1))
        logp1 = F.log_softmax(logits1, dim=-1)
        logp2 = F.log_softmax(logits2, dim=-1)
        p1 = torch.softmax(logits1, dim=-1)
        p2 = torch.softmax(logits2, dim=-1)
        kl1 = F.kl_div(logp1, p2, reduction="batchmean")
        kl2 = F.kl_div(logp2, p1, reduction="batchmean")
        kl = (kl1 + kl2) / 2.0
        loss = (loss1 + loss2) / 2.0 + alpha * kl
    return loss

Rdrop loss.

Classes

class MethodTrainer (method, rdrop_alpha, r3f_epsilon, is_regression, *args, **kwargs)

Expand source code

class MethodTrainer(Trainer):
    def __init__(self, method, rdrop_alpha, r3f_epsilon, is_regression, *args, **kwargs) -> None:
        """Execute __init__ operation."""
        super().__init__(*args, **kwargs)
        self.method = method
        self.rdrop_alpha = rdrop_alpha
        self.r3f_epsilon = r3f_epsilon
        self.is_regression = is_regression

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs) -> Any:
        """Compute and return loss."""
        if "rdrop" in self.method.lower():
            outputs1 = model(**inputs)
            outputs2 = model(**inputs)
            loss = rdrop_loss(
                outputs1.logits,
                outputs2.logits,
                inputs["labels"],
                self.is_regression,
                self.rdrop_alpha,
            )
            return (loss, outputs1) if return_outputs else loss
        elif "r3f" in self.method.lower():
            loss, logits = r3f_forward(
                model,
                input_ids=inputs["input_ids"],
                attention_mask=inputs.get("attention_mask"),
                token_type_ids=inputs.get("token_type_ids"),
                labels=inputs["labels"],
                epsilon=self.r3f_epsilon,
                is_regression=self.is_regression,
            )
            return (loss, logits) if return_outputs else loss
        return super().compute_loss(model, inputs, return_outputs, **kwargs)

Trainer is a simple but feature-complete training and eval loop for PyTorch, optimized for 🤗 Transformers.

Args -----= model ([PreTrainedModel] or torch.nn.Module, optional): The model to train, evaluate or use for predictions. If not provided, a model_init must be passed.

<Tip>

[`Trainer`] is optimized to work with the [`PreTrainedModel`] provided by the library. You can still use
your own models defined as <code>torch.nn.Module</code> as long as they work the same way as the 🤗 Transformers
models.

</Tip>

args ([TrainingArguments], optional): The arguments to tweak for training. Will default to a basic instance of [TrainingArguments] with the output_dir set to a directory named tmp_trainer in the current directory if not provided. data_collator (DataCollator, optional): The function to use to form a batch from a list of elements of train_dataset or eval_dataset. Will default to [default_data_collator] if no processing_class is provided, an instance of [DataCollatorWithPadding] otherwise if the processing_class is a feature extractor or tokenizer. train_dataset (torch.utils.data.Dataset | torch.utils.data.IterableDataset | datasets.Dataset, optional): The dataset to use for training. If it is a [~datasets.Dataset], columns not accepted by the model.forward() method are automatically removed.

Note that if it's a <code>torch.utils.data.IterableDataset</code> with some randomization and you are training in a
distributed fashion, your iterable dataset should either use a internal attribute <code>generator</code> that is a
<code>torch.Generator</code> for the randomization that must be identical on all processes (and the Trainer will
manually set the seed of this <code>generator</code> at each epoch) or have a <code>set\_epoch()</code> method that internally
sets the seed of the RNGs used.

eval_dataset (torch.utils.data.Dataset | dict[str, torch.utils.data.Dataset] | datasets.Dataset, optional): The dataset to use for evaluation. If it is a [~datasets.Dataset], columns not accepted by the model.forward() method are automatically removed. If it is a dictionary, it will evaluate on each dataset prepending the dictionary key to the metric name. processing_class (PreTrainedTokenizerBase or BaseImageProcessor or FeatureExtractionMixin or ProcessorMixin, optional): Processing class used to process the data. If provided, will be used to automatically process the inputs for the model, and it will be saved along the model to make it easier to rerun an interrupted training or reuse the fine-tuned model. model_init (Callable[[], PreTrainedModel], optional): A function that instantiates the model to be used. If provided, each call to [~Trainer.train] will start from a new instance of the model as given by this function.

The function may have zero argument, or a single one containing the optuna/Ray Tune trial object, to
be able to choose different architectures according to hyperparameters (such as layer count, sizes of
inner layers, dropout probabilities etc).

compute_loss_func (Callable, optional): A function that accepts the raw model outputs, labels, and the number of items in the entire accumulated batch (batch_size * gradient_accumulation_steps) and returns the loss. For example, see the default loss function used by [Trainer]. compute_metrics (Callable[[EvalPrediction], Dict], optional): The function that will be used to compute metrics at evaluation. Must take a [EvalPrediction] and return a dictionary string to metric values. Note When passing TrainingArgs with batch_eval_metrics set to True, your compute_metrics function must take a boolean compute_result argument. This will be triggered after the last eval batch to signal that the function needs to calculate and return the global summary statistics rather than accumulating the batch-level statistics callbacks (List of [TrainerCallback], optional): A list of callbacks to customize the training loop. Will add those to the list of default callbacks detailed in here.

If you want to remove one of the default callbacks used, use the [`Trainer.remove_callback`] method.

optimizers (tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR], optional, defaults to (None, None)): A tuple containing the optimizer and the scheduler to use. Will default to an instance of [AdamW] on your model and a scheduler given by [get_linear_schedule_with_warmup] controlled by args. optimizer_cls_and_kwargs (tuple[Type[torch.optim.Optimizer], dict[str, Any]], optional): A tuple containing the optimizer class and keyword arguments to use. Overrides optim and optim_args in args. Incompatible with the optimizers argument.

Unlike <code>optimizers</code>, this argument avoids the need to place model parameters on the correct devices before initializing the Trainer.

preprocess_logits_for_metrics (Callable[[torch.Tensor, torch.Tensor], torch.Tensor], optional): A function that preprocess the logits right before caching them at each evaluation step. Must take two tensors, the logits and the labels, and return the logits once processed as desired. The modifications made by this function will be reflected in the predictions received by compute_metrics.

Note that the labels (second parameter) will be <code>None</code> if the dataset does not have them.

Important attributes:

- **model** -- Always points to the core model. If using a transformers model, it will be a [`PreTrainedModel`]
  subclass.
- **model_wrapped** -- Always points to the most external model in case one or more other modules wrap the
  original model. This is the model that should be used for the forward pass. For example, under <code>DeepSpeed</code>,
  the inner model is wrapped in <code>DeepSpeed</code> and then again in <code>torch.nn.DistributedDataParallel</code>. If the inner
  model hasn't been wrapped, then <code>self.model\_wrapped</code> is the same as <code>self.model</code>.
- **is_model_parallel** -- Whether or not a model has been switched to a model parallel mode (different from
  data parallelism, this means some of the model layers are split on different GPUs).
- **place_model_on_device** -- Whether or not to automatically place the model on the device. Defaults to
  <code>True</code> unless model parallel, DeepSpeed, FSDP, full fp16/bf16 eval, or SageMaker MP is active. Can be
  overridden by subclassing <code>TrainingArguments</code> and overriding the <code>place\_model\_on\_device</code> property.
- **is_in_train** -- Whether or not a model is currently running <code>train</code> (e.g. when <code>evaluate</code> is called while
  in <code>train</code>)

Execute init operation.

Ancestors

• transformers.trainer.Trainer

Methods

def compute_loss(self, model, inputs, return_outputs=False, **kwargs) ‑> Any

Expand source code

def compute_loss(self, model, inputs, return_outputs=False, **kwargs) -> Any:
    """Compute and return loss."""
    if "rdrop" in self.method.lower():
        outputs1 = model(**inputs)
        outputs2 = model(**inputs)
        loss = rdrop_loss(
            outputs1.logits,
            outputs2.logits,
            inputs["labels"],
            self.is_regression,
            self.rdrop_alpha,
        )
        return (loss, outputs1) if return_outputs else loss
    elif "r3f" in self.method.lower():
        loss, logits = r3f_forward(
            model,
            input_ids=inputs["input_ids"],
            attention_mask=inputs.get("attention_mask"),
            token_type_ids=inputs.get("token_type_ids"),
            labels=inputs["labels"],
            epsilon=self.r3f_epsilon,
            is_regression=self.is_regression,
        )
        return (loss, logits) if return_outputs else loss
    return super().compute_loss(model, inputs, return_outputs, **kwargs)

Compute and return loss.