GridSearchCV with KerasClassifier causes Could not pickle the task to send it to the workers with an adapted Keras layer - tensorflow

Question

Question
GridSearchCV via KerasClassifier causes the error when Keras Normalization has been adapted to data. Without the adapted Normalization, it works. The reason why using Normalization is because it gave better result than simply divide by 255.0.
PicklingError: Could not pickle the task to send it to the workers.
Workaround
By setting n_jobs=1 not to multi-thread, it works but perhaps not much use to run single thread.
Environment
Python 3.9.13
TensorFlow version: 2.10.0
Eager execution is: True
Keras version: 2.10.0
sklearn version: 1.1.3
Code
import numpy as np
import tensorflow as tf
from keras.layers import (
Dense,
Flatten,
Normalization,
Conv2D,
MaxPooling2D,
)
from keras.models import (
Sequential
)
from scikeras.wrappers import (
KerasClassifier,
)
from sklearn.model_selection import (
GridSearchCV
)
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
# max_value = float(np.max(x_train))
# x_train, x_test = x_train/max_value, x_test/max_value
input_shape = x_train[0].shape
number_of_classes = 10
# Data Normalization
normalization = Normalization(
name="norm",
input_shape=input_shape, # (32, 32, 3)
axis=-1 # Regard each pixel as a feature
)
normalization.adapt(x_train)
def create_model():
model = Sequential([
# Without the adapted Normalization layer, it works.
normalization,
Conv2D(
name="conv",
filters=32,
kernel_size=(3, 3),
strides=(1, 1),
padding="same",
activation='relu',
input_shape=input_shape
),
MaxPooling2D(
name="maxpool",
pool_size=(2, 2)
),
Flatten(),
Dense(
name="full",
units=100,
activation="relu"
),
Dense(
name="label",
units=number_of_classes,
activation="softmax"
)
])
model.compile(
loss=tf.keras.losses.sparse_categorical_crossentropy,
optimizer='adam',
metrics=['accuracy']
)
return model
model = KerasClassifier(model=create_model, verbose=2)
batch_size = [32]
epochs = [2, 3]
param_grid = dict(batch_size=batch_size, epochs=epochs)
grid = GridSearchCV(estimator=model, param_grid=param_grid, n_jobs=-1, cv=3)
grid_result = grid.fit(x_train, y_train)
Log
The above exception was the direct cause of the following exception:
PicklingError Traceback (most recent call last)
Cell In [28], line 7
4 param_grid = dict(batch_size=batch_size, epochs=epochs)
6 grid = GridSearchCV(estimator=model, param_grid=param_grid, n_jobs=-1, cv=3)
----> 7 grid_result = grid.fit(x_train, y_train)
File ~/venv/ml/lib/python3.9/site-packages/sklearn/model_selection/_search.py:875, in BaseSearchCV.fit(self, X, y, groups, **fit_params)
869 results = self._format_results(
870 all_candidate_params, n_splits, all_out, all_more_results
871 )
873 return results
--> 875 self._run_search(evaluate_candidates)
877 # multimetric is determined here because in the case of a callable
878 # self.scoring the return type is only known after calling
879 first_test_score = all_out[0]["test_scores"]
File ~/venv/ml/lib/python3.9/site-packages/sklearn/model_selection/_search.py:1379, in GridSearchCV._run_search(self, evaluate_candidates)
1377 def _run_search(self, evaluate_candidates):
1378 """Search all candidates in param_grid"""
-> 1379 evaluate_candidates(ParameterGrid(self.param_grid))
File ~/venv/ml/lib/python3.9/site-packages/sklearn/model_selection/_search.py:822, in BaseSearchCV.fit.<locals>.evaluate_candidates(candidate_params, cv, more_results)
814 if self.verbose > 0:
815 print(
816 "Fitting {0} folds for each of {1} candidates,"
817 " totalling {2} fits".format(
818 n_splits, n_candidates, n_candidates * n_splits
819 )
820 )
--> 822 out = parallel(
823 delayed(_fit_and_score)(
824 clone(base_estimator),
825 X,
826 y,
827 train=train,
828 test=test,
829 parameters=parameters,
830 split_progress=(split_idx, n_splits),
831 candidate_progress=(cand_idx, n_candidates),
832 **fit_and_score_kwargs,
833 )
834 for (cand_idx, parameters), (split_idx, (train, test)) in product(
835 enumerate(candidate_params), enumerate(cv.split(X, y, groups))
836 )
837 )
839 if len(out) < 1:
840 raise ValueError(
841 "No fits were performed. "
842 "Was the CV iterator empty? "
843 "Were there no candidates?"
844 )
File ~/venv/ml/lib/python3.9/site-packages/joblib/parallel.py:1098, in Parallel.__call__(self, iterable)
1095 self._iterating = False
1097 with self._backend.retrieval_context():
-> 1098 self.retrieve()
1099 # Make sure that we get a last message telling us we are done
1100 elapsed_time = time.time() - self._start_time
File ~/venv/ml/lib/python3.9/site-packages/joblib/parallel.py:975, in Parallel.retrieve(self)
973 try:
974 if getattr(self._backend, 'supports_timeout', False):
--> 975 self._output.extend(job.get(timeout=self.timeout))
976 else:
977 self._output.extend(job.get())
File ~/venv/ml/lib/python3.9/site-packages/joblib/_parallel_backends.py:567, in LokyBackend.wrap_future_result(future, timeout)
564 """Wrapper for Future.result to implement the same behaviour as
565 AsyncResults.get from multiprocessing."""
566 try:
--> 567 return future.result(timeout=timeout)
568 except CfTimeoutError as e:
569 raise TimeoutError from e
File /Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/concurrent/futures/_base.py:446, in Future.result(self, timeout)
444 raise CancelledError()
445 elif self._state == FINISHED:
--> 446 return self.__get_result()
447 else:
448 raise TimeoutError()
File /Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/concurrent/futures/_base.py:391, in Future.__get_result(self)
389 if self._exception:
390 try:
--> 391 raise self._exception
392 finally:
393 # Break a reference cycle with the exception in self._exception
394 self = None
PicklingError: Could not pickle the task to send it to the workers.
Research
Keras KerasClassifier gridsearch TypeError: can't pickle _thread.lock objects told Keras did not support pickle was the cause. However, as the code works if the adapted Normalization is not used, not relevant.
GPU can cause the issue but there is no GPU in my environment.
References
SciKeras Basic usage¶
How to Grid Search Hyperparameters for Deep Learning Models in Python with Keras