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__)