Refactored hyperparameter tuning in train_utils.py

deepof/train_utils.py

@@ -13,7 +13,6 @@ from kerastuner import BayesianOptimization
 from typing import Tuple, Union, Any, List
 import deepof.hypermodels
 import deepof.model_utils
-import keras
 import numpy as np
 import os
 import pickle
@@ -84,7 +83,7 @@ def get_callbacks(
     )
 
     log_dir = os.path.abspath("logs/fit/{}".format(run_ID))
-    tensorboard_callback = keras.callbacks.TensorBoard(
+    tensorboard_callback = tf.keras.callbacks.TensorBoard(
         log_dir=log_dir, histogram_freq=1, profile_batch=2,
     )
 

tests/test_train_utils.py

@@ -14,7 +14,6 @@ from hypothesis import strategies as st
 from hypothesis.extra.numpy import arrays
 import deepof.model_utils
 import deepof.train_utils
-import keras
 import os
 import tensorflow as tf
 
@@ -64,7 +63,7 @@ def test_get_callbacks(
         X_train, batch_size, variational, predictor, k, loss, kl_wu, mmd_wu,
     )
     assert type(runID) == str
-    assert type(tbc) == keras.callbacks.tensorboard_v2.TensorBoard
+    assert type(tbc) == tf.keras.callbacks.TensorBoard
     assert type(cpc) == tf.python.keras.callbacks.ModelCheckpoint
     assert type(cycle1c) == deepof.model_utils.one_cycle_scheduler