Finegrain/refiners
1
0
Fork 0
mirror of https://github.com/finegrain-ai/refiners.git synced 2024-11-24 15:18:46 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

add before/after init callback to trainer

Browse source
This commit is contained in:
limiteinductive 2023-12-11 14:33:06 +01:00 committed by Benjamin Trom
parent 42a0fc4aa0
commit 7992258dd2
2 changed files with 8 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
src/refiners/training_utils/callback.py
View file

@ -43,6 +43,12 @@ T = TypeVar("T")
class Callback(Generic[T]): class Callback(Generic[T]):
def on_init_begin(self, trainer: T) -> None:
...
def on_init_end(self, trainer: T) -> None:
...
def on_train_begin(self, trainer: T) -> None: def on_train_begin(self, trainer: T) -> None:
... ...

2
src/refiners/training_utils/trainer.py
View file

@ -278,10 +278,12 @@ class Trainer(Generic[ConfigType, Batch]):
) )
self.callbacks = callbacks or [] self.callbacks = callbacks or []
self.callbacks += self.default_callbacks() self.callbacks += self.default_callbacks()
self._call_callbacks(event_name="on_init_begin")
self.load_wandb() self.load_wandb()
self.load_models() self.load_models()
self.prepare_models() self.prepare_models()
self.prepare_checkpointing() self.prepare_checkpointing()
self._call_callbacks(event_name="on_init_end")
def default_callbacks(self) -> list[Callback[Any]]: def default_callbacks(self) -> list[Callback[Any]]:
return [ return [