Pytorch Logging

Import necessary packages

from katonic.ml import MyClient
import os
import numpy as np

import torch
from torch import nn

from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import train_test_split
import pandas as pd

Create client object using MyClient

myclient = MyClient()
mlflow = myclient.mlflow
client = myclient.client

Define evalution function

def eval_metrics(actual, pred):
    rmse = np.sqrt(mean_squared_error(actual, pred))
    mae = mean_absolute_error(actual, pred)
    r2 = r2_score(actual, pred)
    return rmse, mae, r2

Implement model training and logging steps

mlflow.pytorch.autolog()
net = nn.Linear(11, 1)
loss_function = nn.L1Loss()
optimizer = torch.optim.Adam(net.parameters(), lr=1e-4)

csv_url = (
    "https://raw.githubusercontent.com/mlflow/mlflow/master/tests/datasets/winequality-red.csv"
)
data = pd.read_csv(csv_url, sep=";")
train, test = train_test_split(data)

# The predicted column is "quality" which is a scalar from [3, 9]
train_x = train.drop(["quality"], axis=1)
test_x = test.drop(["quality"], axis=1)
train_y = train[["quality"]]
test_y = test[["quality"]]

train_x = torch.tensor(train_x.values.astype(np.float32))
train_y = torch.tensor(train_y.values.astype(np.float32))
test_x = torch.tensor(test_x.values.astype(np.float32))
#test_y = torch.tensor(test_y.values.astype(np.float32))

epochs = 5
for epoch in range(epochs):
    optimizer.zero_grad()
    outputs = net(train_x)
    loss = loss_function(outputs, train_y)
    loss.backward()
    optimizer.step()

exp_name = "mlflow-test-torch"
mlflow.set_experiment(exp_name)
exp_details = mlflow.get_experiment_by_name(exp_name)
with mlflow.start_run(run_name=exp_name):

    predictions = pd.DataFrame(net(test_x).detach().numpy())
    (rmse, mae, r2) = eval_metrics(test_y, predictions)

    mlflow.log_metric("rmse", rmse)
    mlflow.log_metric("r2", r2)
    mlflow.log_metric("mae", mae)    
    mlflow.pytorch.log_model(net, "model")