mlflow fix3
This commit is contained in:
parent
5c28e58e2b
commit
74f9d0ae01
12
train.py
12
train.py
@ -103,14 +103,14 @@ with mlflow.start_run() as run:
|
|||||||
optimizer.step()
|
optimizer.step()
|
||||||
|
|
||||||
torch.save(model,"classificationn_model.pt")
|
torch.save(model,"classificationn_model.pt")
|
||||||
signature = infer_signature(X_train.numpy(), model(X_train).detach().numpy())
|
signature = infer_signature(X_train.numpy(), model(X_train).detach().numpy())
|
||||||
input_example = {"input": X_train[0].numpy().tolist()}
|
input_example = {"input": X_train[0].numpy().tolist()}
|
||||||
|
|
||||||
# Log model
|
# Log model
|
||||||
tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
|
tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
|
||||||
if tracking_url_type_store != "file":
|
if tracking_url_type_store != "file":
|
||||||
mlflow.pytorch.log_model(model, "model", signature=signature, input_example=input_example, registered_model_name="ClassificationModel")
|
mlflow.pytorch.log_model(model, "model", signature=signature, input_example=input_example, registered_model_name="ClassificationModel")
|
||||||
else:
|
else:
|
||||||
mlflow.pytorch.log_model(model, "model", signature=signature, input_example=input_example)
|
mlflow.pytorch.log_model(model, "model", signature=signature, input_example=input_example)
|
||||||
|
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user