Zaktualizuj 'train-mlflow.py'

This commit is contained in:
Kornelia Girejko 2022-05-15 14:24:42 +02:00
parent 9122c49284
commit 154c6ab254

View File

@ -70,7 +70,6 @@ X_testing = torch.from_numpy(X_testing.astype(np.float32))
y_training = torch.from_numpy(y_training.astype(np.float32))
y_testing = torch.from_numpy(y_testing.astype(np.float32))
model = Model()
criterion = nn.BCELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
@ -100,8 +99,13 @@ def my_main(epochs):
#result = open("output",'w+')
#result.write(f'{y_predicted}')
input_example = X_testing[0]
siganture = mlflow.models.signature.infer_signature(X_training, model.predict(X_training))
tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
mlflow.pytorch.log_model(model, "model", registered_model_name="s478815", signature=siganture, input_example=input_example)
torch.save(model, "modelS.pkl")
torch.save(model, "model.pkl")
rmse = mean_squared_error(y_testing, y_predicted)
#print(rmse)