diff --git a/docs/source/main_classes/callback.rst b/docs/source/main_classes/callback.rst index 16b1318b717f..f146244c1fd9 100644 --- a/docs/source/main_classes/callback.rst +++ b/docs/source/main_classes/callback.rst @@ -13,7 +13,7 @@ subclass :class:`~transformers.Trainer` and override the methods you need (see : By default a :class:`~transformers.Trainer` will use the following callbacks: - :class:`~transformers.DefaultFlowCallback` which handles the default behavior for logging, saving and evaluation. -- :class:`~transformers.PrinterCallback` or :class:`~transformers.ProrgressCallback` to display progress and print the +- :class:`~transformers.PrinterCallback` or :class:`~transformers.ProgressCallback` to display progress and print the logs (the first one is used if you deactivate tqdm through the :class:`~transformers.TrainingArguments`, otherwise it's the second one). - :class:`~transformers.integrations.TensorBoardCallback` if tensorboard is accessible (either through PyTorch >= 1.4 @@ -21,6 +21,8 @@ By default a :class:`~transformers.Trainer` will use the following callbacks: - :class:`~transformers.integrations.WandbCallback` if `wandb `__ is installed. - :class:`~transformers.integrations.CometCallback` if `comet_ml `__ is installed. - :class:`~transformers.integrations.MLflowCallback` if `mlflow `__ is installed. +- :class:`~transformers.integrations.AzureMLCallback` if `azureml-sdk `__ is + installed. The main class that implements callbacks is :class:`~transformers.TrainerCallback`. It gets the :class:`~transformers.TrainingArguments` used to instantiate the :class:`~transformers.Trainer`, can access that @@ -50,6 +52,7 @@ Here is the list of the available :class:`~transformers.TrainerCallback` in the .. autoclass:: transformers.integrations.MLflowCallback :members: setup +.. autoclass:: transformers.integrations.AzureMLCallback TrainerCallback ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/src/transformers/integrations.py b/src/transformers/integrations.py index 3925c1838ebf..b734a565e0e1 100644 --- a/src/transformers/integrations.py +++ b/src/transformers/integrations.py @@ -53,6 +53,13 @@ except ImportError: _has_tensorboard = False +try: + from azureml.core.run import Run # noqa: F401 + + _has_azureml = True +except ImportError: + _has_azureml = False + try: import mlflow # noqa: F401 @@ -60,7 +67,6 @@ except ImportError: _has_mlflow = False - # No transformer imports above this point from .file_utils import is_torch_tpu_available @@ -93,6 +99,10 @@ def is_ray_available(): return _has_ray +def is_azureml_available(): + return _has_azureml + + def is_mlflow_available(): return _has_mlflow @@ -420,6 +430,27 @@ def on_log(self, args, state, control, model=None, logs=None, **kwargs): experiment._log_metrics(logs, step=state.global_step, epoch=state.epoch, framework="transformers") +class AzureMLCallback(TrainerCallback): + """ + A :class:`~transformers.TrainerCallback` that sends the logs to `AzureML + `__. + """ + + def __init__(self, azureml_run=None): + assert _has_azureml, "AzureMLCallback requires azureml to be installed. Run `pip install azureml-sdk`." + self.azureml_run = azureml_run + + def on_init_end(self, args, state, control, **kwargs): + if self.azureml_run is None and state.is_world_process_zero: + self.azureml_run = Run.get_context() + + def on_log(self, args, state, control, logs=None, **kwargs): + if self.azureml_run: + for k, v in logs.items(): + if isinstance(v, (int, float)): + self.azureml_run.log(k, v, description=k) + + class MLflowCallback(TrainerCallback): """ A :class:`~transformers.TrainerCallback` that sends the logs to `MLflow `__. diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index f032d4e07563..0eae2321f156 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -40,6 +40,7 @@ from .integrations import ( default_hp_search_backend, hp_params, + is_azureml_available, is_comet_available, is_mlflow_available, is_optuna_available, @@ -151,6 +152,11 @@ if is_ray_available(): from ray import tune +if is_azureml_available(): + from .integrations import AzureMLCallback + + DEFAULT_CALLBACKS.append(AzureMLCallback) + logger = logging.get_logger(__name__)