MLFlow sklearn Model Register的两种写法

MLFlow sklearn Model Register的两种写法

这里介绍MLFlow Model Register的两种写法：

• mlflow.sklearn.log_model
• mlflow.register_model

---

---

比如我们已经导入了数据，并且经过了数据拆分，归一化等预处理：

X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=1)
# 去均值和方差的归一化
sc = StandardScaler()
X_train = sc.fit_transform(X_train)
X_val = sc.transform(X_val)

然后我们新建了一个experiment，以及run

myexperiment = Experiment(workspace, "RFTest1")
mlflow.set_experiment("RFTest1")
# 在此experiment下新建一个run
with mlflow.start_run() as new_run:
      mlflow.log_param("dataset_name", datasetTrn.name)
      mlflow.log_param("dataset_version", datasetTrn.version)

然后我们进行训练以及验证：

rf = RandomForestClassifier(max_depth=10, random_state=0, n_estimators=100)
parameters = {'kernel':('linear', 'rbf'), 'C':[1, 10]}
rf.fit(X_train, y_train)
predicted_rf = rf.predict(X_val)
acc = accuracy_score(y_val, predicted_rf)

最后我们希望将这个模型打包注册，这样我们需要的时候，就可以从需要的版本中把当时训练的模型提取出来继续使用。这里有两个方式进行模型注册：

第一种方式是mlflow.sklearn.log_model：

mlflow.sklearn.log_model(
            sk_model = self.rf, 
            artifact_path = 'outputs/rf.onnx',
            registered_model_name = "sk-learn-random-forest-reg-model" )

第二种方式是mlflow.register_model：

experiment_name='RFTest1'
current_experiment=mlflow.get_experiment_by_name(experiment_name)
runs = mlflow.search_runs(experiment_ids=current_experiment.experiment_id, run_view_type=ViewType.ALL)
run_id = runs.tail(1)["run_id"].tolist()[0]
artifact_path = 'outputs/rf.onnx'
model_uri = 'runs:/{}/{}'.format(run_id, artifact_path) 
registered_model_name = "sk-learn-random-forest-reg-model"
result = mlflow.register_model(
        model_uri = model_uri,
        name = registered_model_name
        )