[logging] Turn off loss logging, while keeping progress bar and logging to third-party application

Feature request

I would like to add a training argument to the TrainingArguments class to turn off the loss logging to stdout while keeping the progress bar and logging to a third-party application like Weights and Biases.

Motivation

I am working on a project that trains a model with the Trainer class. I need to log the losses at every epoch to Weights and Biases. Here is my code:

training_arguments = TrainingArguments(
    output_dir="./logging_dir",
    num_train_epochs=epochs,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    report_to="wandb",
    logging_strategy="epoch",
)

trainer = Trainer(
    model=model,
    args=training_arguments,
    tokenizer=tokenizer,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

This code is contained in a python file train.py and is launched in the terminal via the command python train.py. However, I don’t want to have the loss printed because it breaks my progress bar as you can see in the image below and I need the progress bar to give some feedback on the training time.

Your contribution

I think that I could add a disable_on_log argument to TrainingArguments. Then in the on_log of the ProgressCallback, a condition should be added like this:

def on_log(self, args, state, control, logs=None, **kwargs):
    if state.is_local_process_zero and self.training_bar is not None and args.disable_on_log:
         _ = logs.pop("total_flos", None)
         self.training_bar.write(str(logs))

[antonioloison]

Yes, I’ve tried that, but this makes my progress bar disappear and keeps the loss logging. I want to do the opposite, keep my progress bar and remove the loss logging.

[github-actions[bot]]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.