Add example for registering callbacks with trainers (#10928)

* Add example for callback registry

Resolves: #9036

* Update callback registry documentation

* Added comments for other ways to register callback
This commit is contained in:
Amala Deshmukh 2021-04-05 12:27:23 -04:00 committed by GitHub
parent 9f4e0c23d6
commit e1c02e018c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -74,6 +74,32 @@ TrainerCallback
.. autoclass:: transformers.TrainerCallback
:members:
Here is an example of how to register a custom callback with the PyTorch :class:`~transformers.Trainer`:
.. code-block:: python
class MyCallback(TrainerCallback):
"A callback that prints a message at the beginning of training"
def on_train_begin(self, args, state, control, **kwargs):
print("Starting training")
trainer = Trainer(
model,
args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
callbacks=[MyCallback] # We can either pass the callback class this way or an instance of it (MyCallback())
)
Another way to register a callback is to call ``trainer.add_callback()`` as follows:
.. code-block:: python
trainer = Trainer(...)
trainer.add_callback(MyCallback)
# Alternatively, we can pass an instance of the callback class
trainer.add_callback(MyCallback())
TrainerState
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~