Log losses/metrics with CustomTrainer(Trainer) class in the same frequency as Trainer, with wandb
Hello !
How could I log losses and metrics with a CustomTrainer(Trainer) class, with the same frequency as the Trainer class ?
The code I’m fine-tuning is logging a total loss with the Trainer class :
`total_loss = mlm_loss + clf_loss`
In the Trainer class, the train loss seems to be computed as following : transformers/src/transformers/trainer.py at a5923d4de7df2fbd1f373dfcfe983216b79b6937 · huggingface/transformers · GitHub
And i need to log the `mlm_loss`, the `perplexity`, and the `clf_loss` during training and evaluation at the same frequency as in the **Trainer** class. I’m using Weights and Biases for logging.
How could I implement that ?
Could you explain the logic of the steps/global_step/accumulation_gradient_step ?
I’ve understood that:
* a **training step** represents 1 forward pass + compute the loss (and sometimes the backward pass), and is called once per batch.
* a **gradient accumulation step** aims to simulate a larger effective batch size. It can run x training steps to accumulate gradients + do one optimizer update.
* a **logging step** : controls the frequency of logging (loss + metrics).
* a **global step** is the total number of training step over all the N epochs.
Are these definition correct ?
Should I also use some callbacks from transformers/src/transformers/trainer_callback.py at v4.53.3 · huggingface/transformers · GitHub ?
I’m getting stuck in the implementation because I haven’t understood the logic.
In my code, I’ve defined
class CustomPreTrainingOutput(ModelOutput):
loss: Optional[torch.FloatTensor] = None
mlm_loss: Optional[torch.FloatTensor] = None
clf_loss: Optional[torch.FloatTensor] = None
logits: Optional[torch.FloatTensor] = None #
mlm_logits: Optional[torch.FloatTensor] = None
clf_logits: Optional[torch.FloatTensor] = None
hidden_states: Optional[Tuple[torch.FloatTensor]] = None
attentions: Optional[Tuple[torch.FloatTensor]] = None
So that the forward method returns this CustomPreTrainingOutput.
In the CustomTrainer(Trainer), I have a `compute_loss` method :
class CustomTrainer(Trainer):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# Initialize accumulators
self._stored_mlm_loss = []
self._stored_clf_loss = []
self._stored_perplexity = []
def compute_loss(self, model, inputs, return_outputs=False):
outputs = model(**inputs)
loss = outputs.loss
mlm_loss = getattr(outputs, 'mlm_loss', torch.tensor(0.0, device=loss.device))
clf_loss = getattr(outputs, 'clf_loss', torch.tensor(0.0, device=loss.device))
try:
perplexity = torch.exp(mlm_loss) if mlm_loss is not None else torch.tensor(0.0, device=loss.device)
except Exception:
perplexity = torch.tensor(0.0, device=loss.device)
# Safe float conversions for logging
loss_value = loss.item()
mlm_loss_value = mlm_loss.detach().cpu().item()
clf_loss_value = clf_loss.detach().cpu().item()
perplexity_value = perplexity.detach().cpu().item()
# Log to console (optional)
logger.info(f"class CustomTrainer(Trainer) [compute_loss] loss: {loss_value:.4f}")
logger.info(f"class CustomTrainer(Trainer) [compute_loss] mlm_loss: {mlm_loss:.4f}")
logger.info(f"class CustomTrainer(Trainer) [compute_loss] clf_loss: {clf_loss:.4f}")
logger.info(f"class CustomTrainer(Trainer) [compute_loss] perplexity: {perplexity_value:.4f}")
# Log to W&B and Trainer metrics
self.log({
"loss": loss_value,
"mlm_loss": mlm_loss_value,
"clf_loss": clf_loss_value,
"perplexity": perplexity_value,
})
# wandb.log(
# {"loss": loss_value,
# "mlm_loss": mlm_loss_value,
# "clf_loss": clf_loss_value,
# "perplexity": perplexity_value
# })
return (loss, outputs) if return_outputs else loss
def evaluate(self, eval_dataset=None, **kwargs) -> Dict[str, float]:
# Reset accumulators before evaluation
self._stored_mlm_loss.clear()
self._stored_clf_loss.clear()
self._stored_perplexity.clear()
eval_output = super().evaluate(eval_dataset=eval_dataset, **kwargs)
if self.state.is_world_process_zero:
self.log({
"eval_mlm_loss": eval_output["eval_mlm_loss"],
"eval_clf_loss": eval_output["eval_clf_loss"],
"eval_perplexity": eval_output["eval_perplexity"],
})
wandb.log({f"eval/{k}": v for k, v in eval_output.items()})
return eval_output
# return super().evaluate(eval_dataset=eval_dataset, **kwargs)
def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
# Ensure inputs are on the same device as model
device = model.device if hasattr(model, 'device') else next(model.parameters()).device
inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
# Standard eval step, get full output
with torch.no_grad():
outputs = model(**inputs)
loss = outputs.loss
# Defensive: handle possible missing attributes
mlm_loss = getattr(outputs, 'mlm_loss', torch.tensor(0.0, device=loss.device))
clf_loss = getattr(outputs, 'clf_loss', torch.tensor(0.0, device=loss.device))
perplexity = getattr(outputs, 'perplexity', torch.exp(loss_mlm))
# Save extra losses for logging (will be averaged later in `compute_metrics`)
self._stored_mlm_loss.append(mlm_loss.detach().cpu().item())
self._stored_clf_loss.append(clf_loss.detach().cpu().item())
self._stored_perplexity.append(perplexity.detach().cpu().item())
return (loss, None, None)
I don’t know how to get the same logging frequency than in the Trainer.
I’ve tried to add callbacks with:
class LossSplitLoggerCallback(TrainerCallback):
def on_log(self, args, state, control, logs=None, **kwargs):
trainer = kwargs["trainer"]
if hasattr(trainer, "_last_mlm_loss"):
logs["mlm_loss"] = trainer._last_lm_loss
if hasattr(trainer, "_last_clf_loss"):
logs["clf_loss"] = trainer._last_clf_loss
And in my main function:
trainer = CustomTrainer(
model=model,
args=training_args,
train_dataset=train_dataset if training_args.do_train else None,
eval_dataset=eval_dataset if training_args.do_eval else None,
tokenizer=tokenizer,
data_collator=data_collator,
compute_metrics=compute_metrics if training_args.do_eval and not is_torch_tpu_available() else None,
preprocess_logits_for_metrics=preprocess_logits_for_metrics if training_args.do_eval and not
is_torch_tpu_available() else None,
callbacks=[LossSplitLoggerCallback()],
)
What is the purpose of _maybe_log_save_evaluate() method ?
Could you give me some guidance, please ?
Thank in advance,
