Small fixes to Keras autologging (mlflow#11063)

Signed-off-by: chenmoneygithub <chen.qian@databricks.com>
StateFarmIns · Feb 20, 2024 · 7d20315 · 7d20315
1 parent 2e4818a
commit 7d20315
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 1 deletion.
diff --git a/mlflow/keras/__init__.py b/mlflow/keras/__init__.py
@@ -21,7 +21,7 @@
         "log_model",
     ]
 else:
-    from mlflow.keras.autolog import autolog
+    from mlflow.keras.autologging import autolog
     from mlflow.keras.callback import MLflowCallback
     from mlflow.keras.load import _load_pyfunc, load_model
     from mlflow.keras.save import (

diff --git a/mlflow/keras/autolog.py → mlflow/keras/autologging.py b/mlflow/keras/autolog.py → mlflow/keras/autologging.py
diff --git a/mlflow/tracking/fluent.py b/mlflow/tracking/fluent.py
@@ -2145,6 +2145,7 @@ def print_auto_logged_info(r):
     # eg: mxnet.gluon is the actual library, mlflow.gluon.autolog is our autolog function for it
     LIBRARY_TO_AUTOLOG_MODULE = {
         "tensorflow": "mlflow.tensorflow",
+        "keras": "mlflow.keras",
         "mxnet.gluon": "mlflow.gluon",
         "xgboost": "mlflow.xgboost",
         "lightgbm": "mlflow.lightgbm",

diff --git a/tests/tracking/fluent/test_fluent_autolog.py b/tests/tracking/fluent/test_fluent_autolog.py
@@ -4,6 +4,7 @@
 from unittest import mock
 
 import fastai
+import keras
 import lightgbm
 import pyspark
 import pyspark.ml
@@ -31,6 +32,7 @@
 
 library_to_mlflow_module_without_spark_datasource = {
     tensorflow: mlflow.tensorflow,
+    keras: mlflow.keras,
     fastai: mlflow.fastai,
     sklearn: mlflow.sklearn,
     xgboost: mlflow.xgboost,