Volume mapping

Dockerfile

@@ -24,7 +24,3 @@ RUN /load_data.sh
 
 RUN chmod +x /grab_avocado.py 
 RUN python3 /grab_avocado.py
-
-# Run the model and train it
-RUN chmod +x /model.py
-RUN python3 /model.py

MLProject

@@ -0,0 +1,12 @@
+name: s478841 regression model
+
+docker_env:
+  image: s478841-image:latest
+
+entry_points:
+  main:
+    parameters:
+      epochs: { type: string, default: "140" }
+      steps: { type: string, default: "10" }
+      save_model: { type: string, default: "--save" }
+    command: "python3 scripts/mlflow_train.py -e {epochs} -s {steps} {save_model}"

jenkins/evaluate.Jenkinsfile

@@ -0,0 +1,51 @@
+pipeline {
+    agent {
+        docker { image 's478841-image:latest' }
+    }
+    parameters {
+        gitParameter branchFilter: 'origin/(.*)', defaultValue: 'develop', name: 'BRANCH_NAME', type:'PT_BRANCH'
+        buildSelector(
+            defaultSelector: upstream(),
+            description: 'Build used for artifacts copying',
+            name:'BUILD_SELECTOR')
+    }
+
+    stages {
+        stage('Copy artifacts') {
+            steps {
+                git branch: "${params.BRANCH_NAME}", url: 'https://git.wmi.amu.edu.pl/s478841/ium_478841.git'
+
+                copyArtifacts filter: 'data/*test*', fingerprintArtifacts: true, projectName: 's478841-create-dataset', selector: buildParameter('BUILD_SELECTOR')
+                copyArtifacts filter: 'data/*model*', fingerprintArtifacts: true, projectName: "s478841-training/${BRANCH_NAME}/", selector: buildParameter('BUILD_SELECTOR')
+                copyArtifacts filter: 'evaluation_results.csv', projectName: "s478841-evaluation/${BRANCH_NAME}/", optional: true
+            }
+        }
+        stage('Evaluate model') {
+            steps {
+                // sh 'chmod +x -R ${env.WORKSPACE}'
+                sh 'python3 scripts/evaluate.py'
+            }
+        }
+        stage('Archive artifacts') {
+            steps {
+                archiveArtifacts artifacts: '*data/evaluation_results.csv', onlyIfSuccessful: true
+                archiveArtifacts artifacts: '*data/plots.png', onlyIfSuccessful: true
+            }
+        }
+    }
+
+    post {
+        success {
+            emailext body: 'SUCCESS', subject: 's478841-evaluation', to: 'e19191c5.uam.onmicrosoft.com@emea.teams.ms'
+        }
+        failure {
+            emailext body: 'FAILURE', subject: 's478841-evaluation', to: 'e19191c5.uam.onmicrosoft.com@emea.teams.ms'
+        }
+        unstable {
+            emailext body: 'UNSTABLE', subject: 's478841-evaluation', to: 'e19191c5.uam.onmicrosoft.com@emea.teams.ms'
+        }
+        changed {
+            emailext body: 'CHANGED', subject: 's478841-evaluation', to: 'e19191c5.uam.onmicrosoft.com@emea.teams.ms'
+        }
+    }
+}

jenkins/training.Jenkinsfile

@@ -0,0 +1,75 @@
+pipeline {
+    agent {
+        docker {
+            image 's478841-image:latest'
+            args '-v /mlruns:/mlruns'
+        }
+    }
+    parameters {
+        string(
+            defaultValue: '140',
+            description: 'epochs number',
+            name: 'epochs'
+        )
+        string(
+            defaultValue: '10',
+            description: 'Number of training steps between loss values logging',
+            name: 'step'
+        )
+        string (
+            defaultValue: '--save',
+            description: 'save model after training',
+            name: 'save_model'
+        )
+    }
+
+    stages {
+        stage('Checkout') {
+            steps {
+                checkout([$class: 'GitSCM', branches: [[name: '*/develop']], extensions: [], userRemoteConfigs: [
+                    [url: 'https://git.wmi.amu.edu.pl/s478841/ium_478841.git']]])
+            }
+        }
+        stage('Copy Artifacts') {
+            steps {
+                copyArtifacts filter: 'data/avocado.data*', fingerprintArtifacts: true, projectName: 's478841-create-dataset', selector: lastSuccessful()
+            }
+        }
+        stage('Model training') {
+            steps {
+                sh "chmod +x -R ${env.WORKSPACE}"
+                sh 'python3 scripts/sacred_train.py -e $epochs -s $step $save_model'
+                sh 'python3 scripts/mlflow_train.py -e $epochs -s $step $save_model'
+            }
+        }
+        stage('Archive artifacts') {
+            steps {
+                archiveArtifacts artifacts: 'mlruns/**', onlyIfSuccessful: true
+                archiveArtifacts artifacts: '*data/predictions.csv', onlyIfSuccessful: true
+                archiveArtifacts artifacts: '*data/model_scripted*', onlyIfSuccessful: true
+                dir('data/training_runs') {
+                    archiveArtifacts artifacts: '**/**', onlyIfSuccessful: true
+                }
+            }
+        }
+    }
+
+    post {
+        success {
+            emailext body: 'SUCCESS', subject: "${env.JOB_NAME}", to: 'e19191c5.uam.onmicrosoft.com@emea.teams.ms'
+            build job: 's478841-evaluation/develop'
+        }
+
+        failure {
+            emailext body: 'FAILURE', subject: "${env.JOB_NAME}", to: 'e19191c5.uam.onmicrosoft.com@emea.teams.ms'
+        }
+
+        unstable {
+            emailext body: 'UNSTABLE', subject: "${env.JOB_NAME}", to: 'e19191c5.uam.onmicrosoft.com@emea.teams.ms'
+        }
+
+        changed {
+            emailext body: 'CHANGED', subject: "${env.JOB_NAME}", to: 'e19191c5.uam.onmicrosoft.com@emea.teams.ms'
+        }
+    }
+}

scripts/evaluate.py

@@ -0,0 +1,42 @@
+from model import AvocadoDataset, evaluate_model
+from torch.utils.data import DataLoader
+from torch.jit import load as load_model
+import pandas as pd
+import matplotlib.pyplot as plt
+import matplotlib
+matplotlib.style.use('ggplot')
+
+
+# * Load the test data
+test_data = DataLoader(AvocadoDataset(
+    './data/avocado.data.test'), batch_size=1, shuffle=False)
+
+# * Load the model
+model = load_model('./data/model_scripted.pt')
+model.eval()
+
+# * Append new inference data
+with open('./data/evaluation_results.csv', 'a+') as f:
+    f.write("{0},{1},{2}\n".format(*evaluate_model(test_data, model)))
+
+# * Load all inference data gathered (till the current one)
+results = pd.read_csv('./data/evaluation_results.csv',
+                      names=['MSE', 'RMSE', 'MAE'])
+
+# * Plot the results
+plt.plot(range(1, len(results)+1), results['MSE'], color='green')
+plt.scatter(range(1, len(results)+1),
+            results['MSE'], label='MSE', color='green', marker='.')
+plt.plot(range(1, len(results)+1), results['RMSE'], color='darkred')
+plt.scatter(range(1, len(results)+1),
+            results['RMSE'], label='RMSE', color='darkorange', marker='.')
+plt.plot(range(1, len(results)+1), results['MAE'], color='blue')
+plt.scatter(range(1, len(results)+1),
+            results['MAE'], label='MAE', color='blue', marker='.')
+plt.xticks(range(1, len(results)+1))
+plt.ylabel('Metric value')
+plt.xlabel('Build number')
+plt.legend()
+
+# * Save figure
+plt.savefig('data/plots.png')

scripts/mlflow_train.py

@@ -0,0 +1,212 @@
+from urllib.parse import urlparse
+import mlflow
+import mlflow.pytorch as model_logger
+
+import argparse
+import pandas as pd
+import numpy as np
+from sklearn.metrics import mean_squared_error, mean_absolute_error
+
+import torch
+from torch import nn
+from torch.utils import data as t_u_data
+
+
+# mlflow.set_tracking_uri("http://localhost:5000")
+mlflow.set_tracking_uri("http://172.17.0.1:5000")
+mlflow.set_experiment("s478841")
+
+
+# * Customized Dataset class (base provided by PyTorch)
+class AvocadoDataset(t_u_data.Dataset):
+    def __init__(self, path: str, target: str = 'AveragePrice'):
+        data = pd.read_csv(path)
+        y = data[target].values.astype('float32')
+        self.y = y.reshape((len(y), 1))
+        self.x_data = data.drop(
+            [target], axis=1).values.astype('float32')
+        self.x_shape = data.drop([target], axis=1).shape
+        # print("Data shape is: ", self.x_data.shape)
+
+    def __len__(self):
+        return len(self.x_data)
+
+    def __getitem__(self, idx):
+        return [self.x_data[idx], self.y[idx]]
+
+    def get_shape(self):
+        return self.x_shape
+
+    def get_splits(self, n_test=0.33):
+        test_size = round(n_test * len(self.x_data))
+        train_size = len(self.x_data) - test_size
+        return t_u_data.random_split(self, [train_size, test_size])
+
+
+class AvocadoRegressor(nn.Module):
+    def __init__(self, input_dim):
+        super(AvocadoRegressor, self).__init__()
+        self.hidden1 = nn.Linear(input_dim, 32)
+        nn.init.xavier_uniform_(self.hidden1.weight)
+        self.act1 = nn.ReLU()
+        self.hidden2 = nn.Linear(32, 8)
+        nn.init.xavier_uniform_(self.hidden2.weight)
+        self.act2 = nn.ReLU()
+        self.hidden3 = nn.Linear(8, 1)
+        nn.init.xavier_uniform_(self.hidden3.weight)
+
+    def forward(self, x):
+        x = self.hidden1(x)
+        x = self.act1(x)
+        x = self.hidden2(x)
+        x = self.act2(x)
+        x = self.hidden3(x)
+        return x
+
+
+def prepare_data(paths):
+    train_dl = t_u_data.DataLoader(AvocadoDataset(
+        paths[0]), batch_size=32, shuffle=True)
+    validate_dl = t_u_data.DataLoader(AvocadoDataset(
+        paths[1]), batch_size=128, shuffle=True)
+    test_dl = t_u_data.DataLoader(AvocadoDataset(
+        paths[2]), batch_size=1, shuffle=False)
+    return train_dl, validate_dl, test_dl
+
+
+def train_model(train_dl, model, epochs, log_step):
+    criterion = nn.MSELoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
+    to_compare = None
+    metrics = None
+    for epoch in range(1, epochs+1):
+        for _, (inputs, targets) in enumerate(train_dl):
+            optimizer.zero_grad()
+            yhat = model(inputs)
+            # * For loss value inspection
+            to_compare = (yhat, targets)
+            loss = criterion(yhat, targets)
+            loss.backward()
+            optimizer.step()
+        if epoch == 1 or (epoch) % log_step == 0:
+            result, target = to_compare[0].detach(
+            ).numpy(), to_compare[1].detach().numpy()
+            metrics = {'train.mse': mean_squared_error(target, result),
+                       'train.mae': mean_absolute_error(target, result),
+                       'train.rmse': mean_squared_error(target, result, squared=False)}
+            # _run.log_scalar("training.RMSE", np.sqrt(mse), epoch)
+            # _run.log_scalar("training.MAE", mae, epoch)
+            # _run.log_scalar('training.MSE', mse, epoch)
+            print(
+                f"Epoch {epoch}\t→\tMSE: {metrics['train.mse']},\tRMSE: {metrics['train.rmse']},\tMAE: {metrics['train.mae']}")
+    return metrics
+
+
+def evaluate_model(test_dl, model):
+    predictions, actuals = list(), list()
+    for _, (inputs, targets) in enumerate(test_dl):
+        yhat = model(inputs)
+        # * retrieve numpy array
+        yhat = yhat.detach().numpy()
+        actual = targets.numpy()
+        actual = actual.reshape((len(actual), 1))
+        # * store predictions
+        predictions.append(yhat)
+        actuals.append(actual)
+    predictions, actuals = np.vstack(predictions), np.vstack(actuals)
+    # * return MSE value
+    mse = mean_squared_error(actuals, predictions)
+    rmse = mean_squared_error(actuals, predictions, squared=False)
+    mae = mean_absolute_error(actuals, predictions)
+    return mse, rmse, mae
+
+
+def predict(row, model):
+    row = row[0].flatten()
+    yhat = model(row)
+    yhat = yhat.detach().numpy()
+    return yhat
+
+
+def main(epochs, save_model, log_step):
+    print(
+        f"Your model will be trained for {epochs} epochs, logging every {log_step} steps. Trained model will {'not ' if save_model else ''}be saved.")
+
+    # * Paths to data
+    avocado_data = ['./data/avocado.data.train',
+                    './data/avocado.data.valid',
+                    './data/avocado.data.test']
+
+    # * Data preparation
+    train_dl, validate_dl, test_dl = prepare_data(paths=avocado_data)
+    print(f"""
+          Train set size: {len(train_dl.dataset)},
+          Validate set size: {len(validate_dl.dataset)}
+          Test set size: {len(test_dl.dataset)}
+          """)
+
+    # * Model definition
+    # ! 66 - in case only regions and type are used (among all the categorical vals)
+    model = AvocadoRegressor(235)
+
+    # * Train model
+    print("Let's start the training, mate!")
+    with mlflow.start_run() as run:
+        print("MLflow run experiment_id: {0}".format(run.info.experiment_id))
+        print("MLflow run artifact_uri: {0}".format(run.info.artifact_uri))
+        metrics = train_model(train_dl=train_dl, model=model,
+                              epochs=epochs, log_step=log_step)
+        mlflow.log_param('epochs', epochs)
+        mlflow.log_metrics(metrics)
+
+        # * Evaluate model
+        val_metrics = {key: val for key, val in zip(
+            ['validate.mse', 'validate.rmse', 'validate.mae'], evaluate_model(validate_dl, model))}
+        print(
+            f"\nEvaluation on VALIDATION set\t→\tMSE: {val_metrics['validate.mse']}, RMSE: {val_metrics['validate.rmse']}, MAE: {val_metrics['validate.mae']}")
+        mlflow.log_metrics(val_metrics)
+
+        test_loss = {key: val for key, val in zip(
+            ['test.mse', 'test.rmse', 'test.mae'], evaluate_model(test_dl, model))}
+        print(
+            f"\nEvaluation on TEST set\t→\tMSE: {test_loss['test.mse']}, RMSE: {test_loss['test.rmse']}, MAE: {test_loss['test.mae']}")
+        mlflow.log_metrics(test_loss)
+
+        # tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
+
+        # if tracking_url_type_store != 'file':
+        #     print('First option')
+        #     model_logger.log_model(
+        #         model, "avocados-model", registered_model_name="AvocadoModel_478841")
+        # else:
+        #     print('Second option')
+        #     model_logger.log_model(model, "model")
+
+
+    # * Save the trained model
+    if save_model:
+        print("Your model has been saved - have a nice day!")
+        scripted_model = torch.jit.script(model)
+        scripted_model.save('./data/model_scripted.pt')
+        # ex.add_artifact('./data/model_scripted.pt')
+
+
+# ex.run()
+if __name__ == '__main__':
+    # * Model parameters
+    parser = argparse.ArgumentParser(description="Script performing logistic regression model training",
+                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument(
+        "-e", "--epochs", default=100, help="Number of epochs the model will be trained for")
+    parser.add_argument(
+        "-s", "--step", default=10, help="Number of steps to repeat logging loss values on")
+    parser.add_argument("--save", action="store_true",
+                        help="Save trained model to file 'trained_model.h5'")
+
+    args = vars(parser.parse_args())
+
+    epochs = int(args['epochs'])
+    save_model = args['save']
+    log_step = int(args['step'])
+
+    main(epochs, save_model, log_step)

scripts/model.py

@@ -1,6 +1,8 @@
+import argparse
+
 import pandas as pd
 import numpy as np
-from sklearn.metrics import mean_squared_error
+from sklearn.metrics import mean_squared_error, mean_absolute_error
 
 import torch
 from torch import nn
@@ -99,7 +101,10 @@ def evaluate_model(test_dl, model):
         actuals.append(actual)
     predictions, actuals = np.vstack(predictions), np.vstack(actuals)
     # * return MSE value
-    return mean_squared_error(actuals, predictions)
+    mse = mean_squared_error(actuals, predictions)
+    rmse = mean_squared_error(actuals, predictions, squared=False)
+    mae = mean_absolute_error(actuals, predictions)
+    return mse, rmse, mae
 
 
 def predict(row, model):
@@ -111,6 +116,21 @@ def predict(row, model):
 
 if __name__ == '__main__':
 
+    # * Model parameters
+    parser = argparse.ArgumentParser(description="Script performing logistic regression model training",
+                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument(
+        "-e", "--epochs", default=100, help="Number of epochs the model will be trained for")
+    parser.add_argument("--save", action="store_true",
+                        help="Save trained model to file 'trained_model.h5'")
+
+    args = vars(parser.parse_args())
+
+    epochs = args['epochs']
+    save_model = args['save']
+    print(
+        f"Your model will be trained for {epochs} epochs. Trained model will {'not ' if save_model else ''}be saved.")
+
     # * Paths to data
     avocado_train = './data/avocado.data.train'
     avocado_valid = './data/avocado.data.valid'
@@ -135,14 +155,21 @@ if __name__ == '__main__':
 
     # * Train model
     print("Let's start the training, mate!")
-    train_model(train_dl, model)
+    train_model(train_dl, model, int(epochs))
 
     # * Evaluate model
-    mse = evaluate_model(validate_dl, model)
-    print(f"\nEvaluation\t→\tMSE: {mse}, RMSE: {np.sqrt(mse)}")
+    mse, rmse, mae = evaluate_model(validate_dl, model)
+    print(f"\nEvaluation\t→\tMSE: {mse}, RMSE: {rmse}, MAE: {mae}")
 
     # * Prediction
     predictions = [(predict(row, model)[0], row[1].item()) for row in test_dl]
     preds_df = pd.DataFrame(predictions, columns=["Prediction", "Target"])
-    print("\nNow predictions - hey ho, let's go!\n", preds_df.head())
+    print("\nNow predictions - hey ho, let's go!\n",
+          preds_df.head(), "\n\n...let's save them\ndum...\ndum...\ndum dum dum...\n\tDUM\n")
     preds_df.to_csv("./data/predictions.csv", index=False)
+
+    # * Save the trained model
+    if save_model:
+        print("Your model has been saved - have a nice day!")
+        scripted_model = torch.jit.script(model)
+        scripted_model.save('./data/model_scripted.pt')

scripts/sacred_train.py

@@ -0,0 +1,206 @@
+from sacred import Experiment
+from sacred.observers import FileStorageObserver, MongoObserver
+
+import argparse
+import pandas as pd
+import numpy as np
+from sklearn.metrics import mean_squared_error, mean_absolute_error
+
+import torch
+from torch import nn
+from torch.utils import data as t_u_data
+
+
+ex = Experiment("478841 sacred_scopes", interactive=True, save_git_info=False)
+ex.observers.append(MongoObserver(
+    url='mongodb://admin:IUM_2021@172.17.0.1:27017', db_name='sacred'))
+ex.observers.append(FileStorageObserver('./data/training_runs'))
+
+
+@ex.config
+def my_config():
+    parser = argparse.ArgumentParser(description="Script performing logistic regression model training",
+                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument(
+        "-e", "--epochs", default=100, help="Number of epochs the model will be trained for")
+    parser.add_argument(
+        "-s", "--step", default=10, help="Number of steps to repeat logging loss values on")
+    parser.add_argument("--save", action="store_true",
+                        help="Save trained model to file 'trained_model.h5'")
+
+    args = vars(parser.parse_args())
+
+    epochs = int(args['epochs'])
+    save_model = args['save']
+    log_step = int(args['step'])
+
+
+# * Customized Dataset class (base provided by PyTorch)
+class AvocadoDataset(t_u_data.Dataset):
+    def __init__(self, path: str, target: str = 'AveragePrice'):
+        data = pd.read_csv(path)
+        y = data[target].values.astype('float32')
+        self.y = y.reshape((len(y), 1))
+        self.x_data = data.drop(
+            [target], axis=1).values.astype('float32')
+        self.x_shape = data.drop([target], axis=1).shape
+        # print("Data shape is: ", self.x_data.shape)
+
+    def __len__(self):
+        return len(self.x_data)
+
+    def __getitem__(self, idx):
+        return [self.x_data[idx], self.y[idx]]
+
+    def get_shape(self):
+        return self.x_shape
+
+    def get_splits(self, n_test=0.33):
+        test_size = round(n_test * len(self.x_data))
+        train_size = len(self.x_data) - test_size
+        return t_u_data.random_split(self, [train_size, test_size])
+
+
+class AvocadoRegressor(nn.Module):
+    def __init__(self, input_dim):
+        super(AvocadoRegressor, self).__init__()
+        self.hidden1 = nn.Linear(input_dim, 32)
+        nn.init.xavier_uniform_(self.hidden1.weight)
+        self.act1 = nn.ReLU()
+        self.hidden2 = nn.Linear(32, 8)
+        nn.init.xavier_uniform_(self.hidden2.weight)
+        self.act2 = nn.ReLU()
+        self.hidden3 = nn.Linear(8, 1)
+        nn.init.xavier_uniform_(self.hidden3.weight)
+
+    def forward(self, x):
+        x = self.hidden1(x)
+        x = self.act1(x)
+        x = self.hidden2(x)
+        x = self.act2(x)
+        x = self.hidden3(x)
+        return x
+
+
+def prepare_data(paths):
+    train_dl = t_u_data.DataLoader(AvocadoDataset(
+        paths[0]), batch_size=32, shuffle=True)
+    validate_dl = t_u_data.DataLoader(AvocadoDataset(
+        paths[1]), batch_size=128, shuffle=True)
+    test_dl = t_u_data.DataLoader(AvocadoDataset(
+        paths[2]), batch_size=1, shuffle=False)
+    return train_dl, validate_dl, test_dl
+
+
+@ex.capture
+def train_model(train_dl, model, epochs, log_step, _run):
+    criterion = nn.MSELoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
+    to_compare = None
+
+    for epoch in range(1, epochs+1):
+        for _, (inputs, targets) in enumerate(train_dl):
+            optimizer.zero_grad()
+            yhat = model(inputs)
+            # * For loss value inspection
+            to_compare = (yhat, targets)
+            loss = criterion(yhat, targets)
+
+            loss.backward()
+            optimizer.step()
+
+        if epoch == 1 or (epoch) % log_step == 0:
+            result, target = to_compare[0].detach(
+            ).numpy(), to_compare[1].detach().numpy()
+            mse = mean_squared_error(target, result)
+            mae = mean_absolute_error(target, result)
+            _run.log_scalar("training.RMSE", np.sqrt(mse), epoch)
+            _run.log_scalar("training.MAE", mae, epoch)
+            _run.log_scalar('training.MSE', mse, epoch)
+
+            print(
+                f"Epoch {epoch}\t→\tMSE: {mse},\tRMSE: {np.sqrt(mse)},\tMAE: {mae}")
+
+
+def evaluate_model(test_dl, model):
+    predictions, actuals = list(), list()
+    for _, (inputs, targets) in enumerate(test_dl):
+        yhat = model(inputs)
+        # * retrieve numpy array
+        yhat = yhat.detach().numpy()
+        actual = targets.numpy()
+        actual = actual.reshape((len(actual), 1))
+        # * store predictions
+        predictions.append(yhat)
+        actuals.append(actual)
+    predictions, actuals = np.vstack(predictions), np.vstack(actuals)
+    # * return MSE value
+    mse = mean_squared_error(actuals, predictions)
+    rmse = mean_squared_error(actuals, predictions, squared=False)
+    mae = mean_absolute_error(actuals, predictions)
+    return mse, rmse, mae
+
+
+def predict(row, model):
+    row = row[0].flatten()
+    yhat = model(row)
+    yhat = yhat.detach().numpy()
+    return yhat
+
+
+@ex.main
+def main(epochs, save_model, log_step, _run):
+    print(
+        f"Your model will be trained for {epochs} epochs. Trained model will {'not ' if save_model else ''}be saved.")
+
+    # * Paths to data
+    avocado_data = ['./data/avocado.data.train',
+                    './data/avocado.data.valid',
+                    './data/avocado.data.test']
+
+    # * Data preparation
+    train_dl, validate_dl, test_dl = prepare_data(paths=avocado_data)
+    print(f"""
+          Train set size: {len(train_dl.dataset)},
+          Validate set size: {len(validate_dl.dataset)}
+          Test set size: {len(test_dl.dataset)}
+          """)
+
+    # * Model definition
+    # ! 66 - in case only regions and type are used (among all the categorical vals)
+    model = AvocadoRegressor(235)
+
+    # * Train model
+    print("Let's start the training, mate!")
+    train_model(train_dl=train_dl, model=model,
+                epochs=epochs, log_step=log_step)
+
+    # * Evaluate model
+    mse, rmse, mae = evaluate_model(validate_dl, model)
+    print(
+        f"\nEvaluation on validation set\t→\tMSE: {mse}, RMSE: {rmse}, MAE: {mae}")
+
+    _run.log_scalar("validation.RMSE", rmse, epochs+1)
+    _run.log_scalar("validation.MAE", mae, epochs+1)
+    _run.log_scalar('validation.MSE', mse, epochs+1)
+
+    # * Prediction
+    predictions = [(predict(row, model)[0], row[1].item()) for row in test_dl]
+    preds_df = pd.DataFrame(predictions, columns=["Prediction", "Target"])
+    test_loss = evaluate_model(test_dl, model)
+
+    print("\nNow predictions - hey ho, let's go!\n", preds_df.head(),
+          f"\nLoss values for test data: \t→\tMSE: {test_loss[0]}, RMSE: {test_loss[1]}, MAE: {test_loss[2]}")
+    print("\n...let's save them\ndum...\ndum...\ndum dum dum...\n\tDUM\n")
+
+    preds_df.to_csv("./data/predictions.csv", index=False)
+
+    # * Save the trained model
+    if save_model:
+        print("Your model has been saved - have a nice day!")
+        scripted_model = torch.jit.script(model)
+        scripted_model.save('./data/model_scripted.pt')
+        ex.add_artifact('./data/model_scripted.pt')
+
+
+ex.run()