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
)