I have an unbalanced data with a couple of classes with relatively smaller sample sizes. I am wondering if there is a way to assign the class weights to BertFor SequenceClassification class, maybe in BertConfig ?, as we can do in nn.CrossEntropyLoss.
from transformers import Trainer
def __init__(self, class_weights, *args, **kwargs):
# You pass the class weights when instantiating the Trainer
self.class_weights = class_weights
def compute_loss(self, model, inputs, return_outputs=False):
How the loss is computed by Trainer. By default, all models return the loss in the first element.
Subclass and override for custom behavior.
if self.label_smoother is not None and "labels" in inputs:
labels = inputs.pop("labels")
labels = None
outputs = model(**inputs)
# Save past state if it exists
# TODO: this needs to be fixed and made cleaner later.
if self.args.past_index >= 0:
self._past = outputs[self.args.past_index]
if labels is not None:
loss = self.label_smoother(outputs, labels)
# We don't use .loss here since the model may return tuples instead of ModelOutput.
# Changes start here
# loss = outputs["loss"] if isinstance(outputs, dict) else outputs
logits = outputs['logits']
criterion = torch.nn.CrossEntropyLoss(weights=self.class_weights)
loss = criterion(logits, inputs['labels'])
# Changes end here
return (loss, outputs) if return_outputs else loss