TrainingArguments cosine_with_restarts: how to define the restart count?

Beginners
1

Fine-tuning SegFormer. The arguments look like:

training_args = TrainingArguments(
    output_dir=outputs_dir,
    overwrite_output_dir=True,
    learning_rate=lr,
    num_train_epochs=epochs,
    lr_scheduler_type="cosine_with_restarts",
    warmup_ratio=1.0 / epochs,
    # ...
)

How do I specify the number of restarts for the scheduler?

1 Like
2

You can see this git issue.

1 Like
3

I don’t know when it was implemented, but you can now set it like this,

lr_scheduler_type="cosine_with_restarts",
lr_scheduler_kwargs={"num_cycles":3},
1 Like