和 joblib 集成

Joblib 是一个集成在 scikit-learn 中用于将机器学习任务并行化的包。在 Mars 中，我们使用 Mars remote 为 joblib 实现了一个后端，用户可以将已有的 scikit-learn 任务使用 Mars 并行化。

你可以使用下面的代码注册和使用 Mars 中的 joblib 后端。

from mars.learn.contrib.joblib import register_mars_backend
register_mars_backend()

完成注册后，你可以使用 Mars 服务地址或者已经创建的 Mars Session 来创建 joblib Mars 后端。如果不指定任何参数，默认或者本地 Mars Session 将会被使用。

import joblib
# create with Mars endpoint
with joblib.parallel_backend('mars', service='http://<host>:<port>'):
    # scikit-learn code
# create with existing Mars session
sess = new_session('http://<host>:<port>')
with joblib.parallel_backend('mars', session=sess):
    # scikit-learn code

一个简单的例子如下。在该例子中，我们使用随机搜索来训练和优化一个 SVM 分类器。你只需要将代码中的 joblib.parallel_backend 替换成你自己的 Mars 服务地址即可执行该代码。

import joblib
import sklearn
from sklearn.datasets import load_digits
from sklearn.model_selection import RandomizedSearchCV
from sklearn.svm import SVC

from mars.learn.contrib.joblib import register_mars_backend
register_mars_backend()

digits = load_digits()
param_space = {
    'C': np.logspace(-6, 6, 30),
    'gamma': np.logspace(-8, 8, 30),
    'tol': np.logspace(-4, -1, 30),
    'class_weight': [None, 'balanced'],
}
model = SVC(kernel='rbf')
search = RandomizedSearchCV(model, param_space, cv=5, n_iter=10, verbose=10)

with joblib.parallel_backend('mars', service='http://<host>:<port>'):
    search.fit(digits.data, digits.target)

需要注意的是，joblib 只应当在小规模数据集上使用，这些数据应当能够存储在单机内存中。对于大规模数据，建议使用 Mars learn 中针对 Mars 数据对象包装的算法。