dnrops
/
transformers
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Support turning off the model uploading in ClearML (#20969) 

* Add support for turning off the model uploading in ClearML

* Add documentation for the CLEARML_LOG_MODEL environment variable

* Adjust new doc addition to the new style

Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>

Co-authored-by: Dudu Lasry <dudu.lasry@viz.ai>
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
This commit is contained in:
Dudu Lasry 2023-01-06 14:22:19 +02:00 committed by GitHub
parent ff8dcb5efa
commit 61e068e5a2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/transformers/integrations.py
View File

@ -1313,6 +1313,8 @@ class ClearMLCallback(TrainerCallback):
ClearML project name.
- **CLEARML_TASK** (`str`, *optional*, defaults to `Trainer`):
ClearML task name.
- **CLEARML_LOG_MODEL** (`bool`, *optional*, defaults to `False`):
Whether to log models as artifacts during training.
"""
def __init__(self):
@ -1326,6 +1328,8 @@ class ClearMLCallback(TrainerCallback):
self._initialized = False
self._clearml_task = None
self._log_model = os.getenv("CLEARML_LOG_MODEL", "FALSE").upper() in ENV_VARS_TRUE_VALUES.union({"TRUE"})
def setup(self, args, state, model, tokenizer, **kwargs):
if self._clearml is None:
return
@ -1403,7 +1407,7 @@ class ClearMLCallback(TrainerCallback):
)
def on_save(self, args, state, control, **kwargs):
if self._clearml_task and state.is_world_process_zero:
if self._log_model and self._clearml_task and state.is_world_process_zero:
ckpt_dir = f"checkpoint-{state.global_step}"
artifact_path = os.path.join(args.output_dir, ckpt_dir)
logger.info(f"Logging checkpoint artifacts in {ckpt_dir}. This may take time.")