Multiple Keras GridSearchCV with multiprocessing: TypeError: cannot pickle '_thread.RLock' object - keras

Question

Edit: I figured that leaving out the early stopping solves the problem, but that is not an option for me. Is there a way to include early stopping in a grid search together with multiprocessing?
I am trying to run GridSeachCV in a Jupyter Notebook for multiple datasets, one after another, using multiprocessing for each grid search. This works fine for the first grid search, but after that I keep getting errors for all other grid searches. I even get errors when I run the same cell again, that worked initially.
import numpy as np
from sklearn.model_selection import GridSearchCV
from tensorflow import keras
from keras.models import Sequential
from keras.layers import Dense
from keras.wrappers.scikit_learn import KerasRegressor
def create_model(n_weights=1000, hidden_layers=1, learning_rate=0.001):
# formulas derived from nWeights = sum (d(l-1)+1)*d(l) for all layers l with output dim d(l)
if hidden_layers == 1: # 100% of neurons in first hidden layer
neurons = [ceil((n_weights - 1) / 3)]
elif hidden_layers == 2: # 70% / 30% split of neurons
x = 1/7 * (np.sqrt(21 * n_weights + 79) - 10)
neurons = list(map(floor,[7/3 * x, x]))
elif hidden_layers == 3: # 50% / 30% / 20% split
x = 1/21 * (np.sqrt(84 * n_weights + 205) - 17)
neurons = list(map(floor, [5/2 * x, 3/2 * x, x]))
else:
raise Exception('Only 1, 2 or 3 layers allowed')
model = Sequential([Dense(neurons[0], activation='relu', input_dim=1)])
for n in neurons[1:]:
model.add(Dense(n, activation='relu'))
model.add(Dense(1))
model.compile(optimizer=keras.optimizers.Adam(learning_rate=learning_rate), loss='mse')
return model
batch_size = [64, 128]
learning_rate = [0.01, 0.001]
hidden_layers = [1, 2, 3]
n_weights = [10, 30, 60]
p_grid = dict(n_weights=n_weights, hidden_layers=hidden_layers, batch_size=batch_size, learning_rate=learning_rate)
earlyStop = keras.callbacks.EarlyStopping(monitor='loss', patience=3, restore_best_weights=True)
epochs = 5000
# first grid search
model1 = KerasRegressor(create_model, epochs=epochs, verbose=0)
grid1 = GridSearchCV(estimator=model, param_grid=p_grid, n_jobs=-1, cv=4, verbose=1)
result1 = grid1.fit(X_train1, y_train1, callbacks=[earlyStop])
# second grid search
model2 = KerasRegressor(create_model, epochs=epochs, verbose=0)
grid2 = GridSearchCV(estimator=model, param_grid=p_grid, n_jobs=-1, cv=4, verbose=1)
result2 = grid2.fit(X_train1, y_train1, callbacks=[earlyStop])
The above code runs two identical grid searches on two different data sets. The first one works fine, but the second fails in this line
grid2 = GridSearchCV(estimator=model, param_grid=p_grid, n_jobs=-1, cv=4, verbose=1)
with the error
_RemoteTraceback Traceback (most recent call last)
_RemoteTraceback:
"""
Traceback (most recent call last):
File "C:\Users\oli-w\anaconda3\envs\tensorflow\lib\site-packages\joblib\externals\loky\backend\queues.py", line 153, in _feed
obj_ = dumps(obj, reducers=reducers)
File "C:\Users\oli-w\anaconda3\envs\tensorflow\lib\site-packages\joblib\externals\loky\backend\reduction.py", line 271, in dumps
dump(obj, buf, reducers=reducers, protocol=protocol)
File "C:\Users\oli-w\anaconda3\envs\tensorflow\lib\site-packages\joblib\externals\loky\backend\reduction.py", line 264, in dump
_LokyPickler(file, reducers=reducers, protocol=protocol).dump(obj)
File "C:\Users\oli-w\anaconda3\envs\tensorflow\lib\site-packages\joblib\externals\cloudpickle\cloudpickle_fast.py", line 602, in dump
return Pickler.dump(self, obj)
TypeError: cannot pickle '_thread.RLock' object
"""
The above exception was the direct cause of the following exception:
PicklingError Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_9700/1799096090.py in <module>
1 # second grid search
2 model2 = KerasRegressor(create_model, epochs=epochs, verbose=0)
----> 3 grid2 = GridSearchCV(estimator=model, param_grid=p_grid, n_jobs=-1, cv=4, verbose=1)
4 result2 = grid2.fit(X_train1, y_train1, callbacks=[earlyStop])
~\anaconda3\envs\tensorflow\lib\site-packages\sklearn\model_selection\_search.py in fit(self, X, y, groups, **fit_params)
889 return results
890
--> 891 self._run_search(evaluate_candidates)
892
893 # multimetric is determined here because in the case of a callable
~\anaconda3\envs\tensorflow\lib\site-packages\sklearn\model_selection\_search.py in _run_search(self, evaluate_candidates)
1390 def _run_search(self, evaluate_candidates):
1391 """Search all candidates in param_grid"""
-> 1392 evaluate_candidates(ParameterGrid(self.param_grid))
1393
1394
~\anaconda3\envs\tensorflow\lib\site-packages\sklearn\model_selection\_search.py in evaluate_candidates(candidate_params, cv, more_results)
836 )
837
--> 838 out = parallel(
839 delayed(_fit_and_score)(
840 clone(base_estimator),
~\anaconda3\envs\tensorflow\lib\site-packages\joblib\parallel.py in __call__(self, iterable)
1054
1055 with self._backend.retrieval_context():
-> 1056 self.retrieve()
1057 # Make sure that we get a last message telling us we are done
1058 elapsed_time = time.time() - self._start_time
~\anaconda3\envs\tensorflow\lib\site-packages\joblib\parallel.py in retrieve(self)
933 try:
934 if getattr(self._backend, 'supports_timeout', False):
--> 935 self._output.extend(job.get(timeout=self.timeout))
936 else:
937 self._output.extend(job.get())
~\anaconda3\envs\tensorflow\lib\site-packages\joblib\_parallel_backends.py in wrap_future_result(future, timeout)
540 AsyncResults.get from multiprocessing."""
541 try:
--> 542 return future.result(timeout=timeout)
543 except CfTimeoutError as e:
544 raise TimeoutError from e
~\anaconda3\envs\tensorflow\lib\concurrent\futures\_base.py in result(self, timeout)
437 raise CancelledError()
438 elif self._state == FINISHED:
--> 439 return self.__get_result()
440 else:
441 raise TimeoutError()
~\anaconda3\envs\tensorflow\lib\concurrent\futures\_base.py in __get_result(self)
386 def __get_result(self):
387 if self._exception:
--> 388 raise self._exception
389 else:
390 return self._result
PicklingError: Could not pickle the task to send it to the workers.
I would understand if the grid search didn't work at all, but it confuses me that it works once and then starts throwing errors. I was only able to run another grid search after I restarted the kernel. What can I do in order to make multiple grid searches possible in one sitting without restarting the kernel?