s487194/ium_487194
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

mlflow fix3

Browse Source
This commit is contained in:
Witold Woch 2023-05-14 23:39:33 +02:00
parent 5c28e58e2b
commit 74f9d0ae01
1 changed files with 8 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
train.py
View File

@ -103,14 +103,14 @@ with mlflow.start_run() as run:
optimizer.step()
torch.save(model,"classificationn_model.pt")
signature = infer_signature(X_train.numpy(), model(X_train).detach().numpy())
input_example = {"input": X_train[0].numpy().tolist()}
signature = infer_signature(X_train.numpy(), model(X_train).detach().numpy())
input_example = {"input": X_train[0].numpy().tolist()}
# Log model
tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
if tracking_url_type_store != "file":
mlflow.pytorch.log_model(model, "model", signature=signature, input_example=input_example, registered_model_name="ClassificationModel")
else:
mlflow.pytorch.log_model(model, "model", signature=signature, input_example=input_example)
# Log model
tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
if tracking_url_type_store != "file":
mlflow.pytorch.log_model(model, "model", signature=signature, input_example=input_example, registered_model_name="ClassificationModel")
else:
mlflow.pytorch.log_model(model, "model", signature=signature, input_example=input_example)