diff --git a/dan/manager/training.py b/dan/manager/training.py index 7b565c41e9465531821f8af5ed739d1794670559..921f1c8a01dc0e228141edf5291bb70b491c004f 100644 --- a/dan/manager/training.py +++ b/dan/manager/training.py @@ -23,6 +23,11 @@ from tqdm import tqdm from dan.manager.metrics import MetricManager from dan.ocr.utils import LM_ind_to_str from dan.schedulers import DropoutScheduler +try: + import mlflow + from dan.mlflow import logging_metrics, logging_tags_metrics +except ImportError: + pass try: import mlflow diff --git a/dan/ocr/document/train.py b/dan/ocr/document/train.py index 93829fe0f402014699f2f236bc3b0d21817a5884..7e2729f18fa75d3a0ba68b6cb2e4041842ed15b4 100644 --- a/dan/ocr/document/train.py +++ b/dan/ocr/document/train.py @@ -28,13 +28,6 @@ try: except ImportError: MLFLOW = False -try: - import mlflow - - MLFLOW = True - logger.info("MLflow Logging available.") -except ImportError: - MLFLOW = False logger = logging.getLogger(__name__) diff --git a/dan/utils.py b/dan/utils.py index 50f7311d602c97e80c308637e1b8cb37d8c90f95..e11c2e85bdc013193be1771be228f10999b81800 100644 --- a/dan/utils.py +++ b/dan/utils.py @@ -23,7 +23,10 @@ class MLflowNotInstalled(Exception): Raised when MLflow logging was requested but the module was not installed """ +<<<<<<< HEAD +======= +>>>>>>> bf9cf44 (check if mlflow is available at runtime) def randint(low, high): """ call torch.randint to preserve random among dataloader workers