"tests/git@developer.sourcefind.cn:wangsen/paddle_dbnet.git" did not exist on "a8a9b2e55124074ef0e75690c8c30a126bb409c3"
Unverified Commit eba418ac authored by Tristan Deleu's avatar Tristan Deleu Committed by GitHub
Browse files

Commit the last step on world_process_zero in WandbCallback (#9805)

* Commit the last step on world_process_zero in WandbCallback

* Use the environment variable WANDB_LOG_MODEL as a default value in WandbCallback
parent 8edc98bb
...@@ -516,6 +516,8 @@ class WandbCallback(TrainerCallback): ...@@ -516,6 +516,8 @@ class WandbCallback(TrainerCallback):
else: else:
self._wandb = wandb self._wandb = wandb
self._initialized = False self._initialized = False
# log outputs
self._log_model = os.getenv("WANDB_LOG_MODEL", "FALSE").upper() in ENV_VARS_TRUE_VALUES.union({"TRUE"})
def setup(self, args, state, model, reinit, **kwargs): def setup(self, args, state, model, reinit, **kwargs):
""" """
...@@ -569,9 +571,6 @@ class WandbCallback(TrainerCallback): ...@@ -569,9 +571,6 @@ class WandbCallback(TrainerCallback):
model, log=os.getenv("WANDB_WATCH", "gradients"), log_freq=max(100, args.logging_steps) model, log=os.getenv("WANDB_WATCH", "gradients"), log_freq=max(100, args.logging_steps)
) )
# log outputs
self._log_model = os.getenv("WANDB_LOG_MODEL", "FALSE").upper() in ENV_VARS_TRUE_VALUES.union({"TRUE"})
def on_train_begin(self, args, state, control, model=None, **kwargs): def on_train_begin(self, args, state, control, model=None, **kwargs):
if self._wandb is None: if self._wandb is None:
return return
...@@ -583,7 +582,8 @@ class WandbCallback(TrainerCallback): ...@@ -583,7 +582,8 @@ class WandbCallback(TrainerCallback):
if self._wandb is None: if self._wandb is None:
return return
# commit last step # commit last step
self._wandb.log({}) if state.is_world_process_zero:
self._wandb.log({})
if self._log_model and self._initialized and state.is_world_process_zero: if self._log_model and self._initialized and state.is_world_process_zero:
from .trainer import Trainer from .trainer import Trainer
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment