How to optimize with optuna an MLP using the Pipeline

Import classes and define paths

from cetaceo.pipeline import Pipeline
from cetaceo.models import MLP
from cetaceo.optimization import OptunaBaseOptimizer
from cetaceo.data import VTUDataset
from cetaceo.evaluators import RegressionEvaluator
from cetaceo.utils import PathManager
from pathlib import Path

import torch
from sklearn.preprocessing import MinMaxScaler

import warnings
warnings.filterwarnings("ignore")

import numpy as np
# since version 1.20.0, numpy has changed the name of the bool type and np.bool is deprecated
# we need to change it to np.bool_ to avoid error with some libraries that still use the old name
np.bool = np.bool_

DATA_DIR = Path.cwd().parent / "sample_data"
CASE_DIR = Path.cwd() / "results"
PathManager.create_directory(CASE_DIR / 'models')
PathManager.create_directory(CASE_DIR / 'hyperparameters')
PathManager.create_directory(CASE_DIR / 'plots')

Define sklearn scalers if needed

Here, we create 2 minmax scalers, one for scaling the inputs, and other for the outputs.

x_scaler = MinMaxScaler()
y_scaler = MinMaxScaler()

Create datasets

For this example, we will use some VTU (VTK like) meshes obtained with CODA. Thus, a VTUDataset is needed. We create one for each dataset split.

The meshes are in separate files

mesh_files = list(DATA_DIR.glob("*.vtu"))
n_test_samples = int(0.4 * len(mesh_files))

train_files = mesh_files[n_test_samples:]
test_files = mesh_files[:n_test_samples // 2]
valid_files = mesh_files[n_test_samples // 2:n_test_samples]

train_dataset = VTUDataset(mesh_files=train_files, x_scaler=x_scaler, y_scaler=y_scaler)
test_dataset = VTUDataset(mesh_files=test_files, x_scaler=x_scaler, y_scaler=y_scaler)
valid_dataset = VTUDataset(mesh_files=valid_files, x_scaler=x_scaler, y_scaler=y_scaler)

x, y = train_dataset[:]
print("Train dataset length :", len(train_dataset))
print("Test dataset length :", len(test_dataset))
print("Valid dataset length :", len(valid_dataset))
print("X, y train shapes:\n", x.shape, y.shape)

Train dataset length : 175224
Test dataset length : 58408
Valid dataset length : 58408
X, y train shapes:
 torch.Size([175224, 2]) torch.Size([175224, 7])

If we need to preprocess the data somehow, we can use the method process_data from the datasets. For instance, we need the velocities from the meshes, but we only have the momentum and the density. We can then apply the following transformation

def process_cell_data(x, y):
    rho = y[..., 0:1]
    # get velocities from momentum and density
    u = y[..., 2:3] / rho
    v = y[..., 3:4] / rho
    p = y[..., 4:5]
    return x, torch.hstack([rho, u, v, p])

train_dataset.process_data(process_function=process_cell_data)
test_dataset.process_data(process_function=process_cell_data)
valid_dataset.process_data(process_function=process_cell_data)

And now we can scale the data

train_dataset.scale_data()
valid_dataset.scale_data()
test_dataset.scale_data()

Evaluator

For this example, we are going to use a RegressionEvaluator, this time no plots will be generated

evaluator = RegressionEvaluator()

Optimization

For optimization, first we need to define a set of parameters.

* If a parameter is not defined, its default value will be taken.

* If a parameter is defined with a single value, that parameter won’t be optimized and that value will be constant during the whole optimization process.

* If a paramer is defined with a tuple with 2 elelements, that parameter will be optimized and the values chosen during the optimization phase will be on the range defined by the tuple.

If we want to store the best parameters obtained during the optimization, a save_dir must be specified on the optimizer constructor. Then, they will be stored in a json file

device = torch.device("cuda:3" if torch.cuda.is_available() else "cpu")
optim_params = {
        "lr": 0.01,  # fixed parameter
        "batch_size": (10, 64),  # optimizable parameter
        "n_layers": (1, 3),
        "hidden_size": 32,
        "print_rate_epoch": 10,
        "epochs": 50,
        "device": device,
    }

optimizer = OptunaBaseOptimizer(
        optimization_params=optim_params, n_trials=5, direction="minimize", save_dir=CASE_DIR / 'hyperparameters'
    )

Run the pipeline

When optimizing, you have to specify the model class to optimize. In this case, we are optimizing an MLP.

pipeline = Pipeline(
        train_dataset=train_dataset,
        test_dataset=test_dataset,
        valid_dataset=valid_dataset,
        model_class=MLP,
        optimizer=optimizer,
        evaluators=[evaluator]
    )
pipeline.run()

[I 2024-09-11 14:54:38,215] A new study created in memory with name: no-name-d6c8ce02-114c-4fd1-8890-c0486091bf01

Epoch 10/50 | Train loss (x1e5) 402.1901
Epoch 20/50 | Train loss (x1e5) 402.2421
Epoch 30/50 | Train loss (x1e5) 402.4459
Epoch 40/50 | Train loss (x1e5) 401.9750
Epoch 50/50 | Train loss (x1e5) 402.0932

[I 2024-09-11 15:12:45,447] Trial 0 finished with value: 0.003009177435055923 and parameters: {'batch_size': 26, 'n_layers': 3}. Best is trial 0 with value: 0.003009177435055923.

Epoch 10/50 | Train loss (x1e5) 334.2382
Epoch 20/50 | Train loss (x1e5) 321.1680
Epoch 30/50 | Train loss (x1e5) 307.4511
Epoch 40/50 | Train loss (x1e5) 299.3205
Epoch 50/50 | Train loss (x1e5) 295.5657

[I 2024-09-11 15:20:02,835] Trial 1 finished with value: 0.0022019218127490216 and parameters: {'batch_size': 47, 'n_layers': 1}. Best is trial 1 with value: 0.0022019218127490216.

Epoch 10/50 | Train loss (x1e5) 322.4274
Epoch 20/50 | Train loss (x1e5) 307.9958
Epoch 30/50 | Train loss (x1e5) 302.4046
Epoch 40/50 | Train loss (x1e5) 295.0588
Epoch 50/50 | Train loss (x1e5) 295.6092

[I 2024-09-11 15:26:58,499] Trial 2 finished with value: 0.002151052706345495 and parameters: {'batch_size': 50, 'n_layers': 1}. Best is trial 2 with value: 0.002151052706345495.

Epoch 10/50 | Train loss (x1e5) 316.7952
Epoch 20/50 | Train loss (x1e5) 305.2839
Epoch 30/50 | Train loss (x1e5) 300.2761
Epoch 40/50 | Train loss (x1e5) 297.5901
Epoch 50/50 | Train loss (x1e5) 294.3364

[I 2024-09-11 15:58:24,977] Trial 3 finished with value: 0.002080027835441825 and parameters: {'batch_size': 13, 'n_layers': 2}. Best is trial 3 with value: 0.002080027835441825.

Epoch 10/50 | Train loss (x1e5) 402.7308
Epoch 20/50 | Train loss (x1e5) 402.7523
Epoch 30/50 | Train loss (x1e5) 402.8952
Epoch 40/50 | Train loss (x1e5) 402.4810
Epoch 50/50 | Train loss (x1e5) 402.5232

[I 2024-09-11 16:16:17,864] Trial 4 finished with value: 0.003013648607862439 and parameters: {'batch_size': 23, 'n_layers': 3}. Best is trial 3 with value: 0.002080027835441825.

Study statistics:
  Number of finished trials:  5
  Number of pruned trials:  0
  Number of completed trials:  5
Best trial:
  Value:  0.002080027835441825
  Params:
    batch_size: 13
    n_layers: 2


Epoch 10/50 | Train loss (x1e5) 404.7949  | Test loss (x1e5) 309.9594
Epoch 20/50 | Train loss (x1e5) 404.7411  | Test loss (x1e5) 307.5095
Epoch 30/50 | Train loss (x1e5) 404.5180  | Test loss (x1e5) 299.7614
Epoch 40/50 | Train loss (x1e5) 404.8046  | Test loss (x1e5) 314.5927
Epoch 50/50 | Train loss (x1e5) 405.0134  | Test loss (x1e5) 299.8414

--------------------------------------------------
Metrics on train data:
--------------------------------------------------

Regression evaluator metrics:
mse: 0.0040
mae: 0.0299
mre: 445.1378%
ae_95: 0.1222
ae_99: 0.3160
r2: 0.9754
l2_error: 0.0834
--------------------------------------------------
Metrics on test data:
--------------------------------------------------

Regression evaluator metrics:
mse: 0.0031
mae: 0.0276
mre: 429.8295%
ae_95: 0.1078
ae_99: 0.2711
r2: 0.9813
l2_error: 0.0726

To save the trained model with the optimized parameters

pipeline.model.save(CASE_DIR / 'models' / 'model.pth')

To load a model, simply use Model.load method

model = MLP.load(CASE_DIR / 'models' / 'model.pth')