| training_args = TrainingArguments( | |
| output_dir="segformer-b0-scene-parse-150", | |
| learning_rate=6e-5, | |
| num_train_epochs=50, | |
| per_device_train_batch_size=2, | |
| per_device_eval_batch_size=2, | |
| save_total_limit=3, | |
| evaluation_strategy="steps", | |
| save_strategy="steps", | |
| save_steps=20, | |
| eval_steps=20, | |
| logging_steps=1, | |
| eval_accumulation_steps=5, | |
| remove_unused_columns=False, | |
| push_to_hub=True, | |
| ) | |
| trainer = Trainer( | |
| model=model, | |
| args=training_args, | |
| train_dataset=train_ds, | |
| eval_dataset=test_ds, | |
| compute_metrics=compute_metrics, | |
| ) | |
| trainer.train() | |
| Once training is completed, share your model to the Hub with the [~transformers.Trainer.push_to_hub] method so everyone can use your model: | |
| trainer.push_to_hub() | |
| If you are unfamiliar with fine-tuning a model with Keras, check out the basic tutorial first! |