Cannot save items in a defaultdict in the scorer of a GridSearchCV with joblib backend

I’m trying to store some information about a training of a custom sklearn-compatible SVM in a shared defautdict in the context of joblib in sklearn GridSearchCV. So I have the following:

    from collections import defaultdict

    from sklearn.model_selection import GridSearchCV
    from sklearn.datasets import make_regression

    X, y = make_regression(n_samples=100, n_features=1, noise=10, bias=2, random_state=123456)

    train_loss_history = defaultdict(lambda: defaultdict(lambda: defaultdict(dict)))

    def dual_r2_scorer(svr, X, y):
        # (f_t - f*) / f*
        f_star = np.full_like(svr.train_loss_history, svr.obj.f_star())
        train_loss_history[svr.dual][svr.kernel][svr.C][svr.epsilon] = np.divide((svr.train_loss_history - f_star),
                                                                                 np.abs(f_star))
        return {'r2': svr.score(X, y),
                'n_iter': svr.optimizer.iter,
                'n_sv': len(svr.support_)}
    

    grid = GridSearchCV(DualSVR(loss=squared_epsilon_insensitive,
                                optimizer=AdaGrad,
                                max_iter=1000,
                                learning_rate=1.),
                        param_grid={'epsilon': [0.1, 0.2, 0.3],
                                    'C': [1, 10, 100]},
                        scoring=dual_r2_scorer,
                        n_jobs=-1,
                        refit='r2',  # refit the best model (wrt r2) on the full development set
                        return_train_score=True,
                        verbose=True).fit(lin_X, lin_y)

But with -1 option for n_jobs, at the end of the execution, the dictionary train_loss_history is empty. For n_jobs=None, it works.

How can I fix it?

[NicolasHug]

I’m glad these docs are helpful ❤️

[NicolasHug]

Here’s a more minimal example:

from collections import defaultdict

from sklearn.model_selection import GridSearchCV
from sklearn.datasets import make_regression
from sklearn.linear_model import LinearRegression

X, y = make_regression()

train_loss_history = defaultdict(str)

def dual_r2_scorer(svr, X, y):
    train_loss_history['hello'] = 12
    return 3


grid = GridSearchCV(LinearRegression(),
                    param_grid={'fit_intercept': [True, False]},
                    scoring=dual_r2_scorer,
                    n_jobs=-1,
                    refit=False,
                    ).fit(X, y)
print(train_loss_history)