from flax import linen as nn import flax import jax import jax.numpy as jnp import numpy as np import argparse import pandas as pd from sklearn.preprocessing import StandardScaler import ast from matplotlib import pyplot as plt from typing import Sequence from flax.training import train_state from flax.training import checkpoints from flax import struct import optax import random import seaborn as sns from flax.training import orbax_utils import shutil from sklearn.model_selection import ParameterGrid from concurrent.futures import ProcessPoolExecutor from itertools import product from flax.training.early_stopping import EarlyStopping import pickle import optuna from typing import Union import sys jax.config.update("jax_enable_x64", True) print(sys.argv) def parse_int_or_none(value): if value != "None": return int(value) else: return None def parse_arguments(): parser = argparse.ArgumentParser() parser.add_argument('--train_data', required=True, type=str, help="Training data file shape") subparsers = parser.add_subparsers(dest="task", help="Task to perform (optimize/train)") # subcommand: optimize optimize_parser = subparsers.add_parser("optimize", help="Optimize task") optimize_parser.add_argument("--number_trials", type=int, required=True, help="Number of trials for optimization") # subcommand: train train_parser = subparsers.add_parser("train", help="Train task") train_parser.add_argument("--LSTM_size", required=True, type=str, help="Hidden size of LSTM layer") train_parser.add_argument("--dense1_size", required=True, type=str, help="Hidden size of 1. Dense layer (None if no layer)") train_parser.add_argument("--dense2_size", required=True, type=str, help="Hidden size of 2. Dense layer (None if no layer)") train_parser.add_argument("--dense3_size", required=True, type=str, help="Hidden size of 3. Dense layer (None if no layer)") train_parser.add_argument("--dropout_rate", type=float,required=True, help="Dropout for training (between 0 and 1)") train_parser.add_argument("--batch_size", type=int, required=True, help="Batch size for training") train_parser.add_argument("--learning_rate", type=float,required=True, help="Learning rate for training") train_parser.add_argument("--number_epochs", type=int, required=True, help="Number of epochs for training") return parser.parse_args() args = parse_arguments() print(args) if args.task == "train": LSTM_size = parse_int_or_none(args.LSTM_size) dense1_size = parse_int_or_none(args.dense1_size) dense2_size = parse_int_or_none(args.dense2_size) dense3_size = parse_int_or_none(args.dense3_size) data_file_shape=args.train_data fold_num=5 def get_batches(df_data, scaler_X, scaler_y, batch_size, train): if train: scaler_X=StandardScaler() scaler_y=StandardScaler() #shuffle df_data=df_data.sample(frac=1) X_df=df_data[['meteo','parameters']] if train: X_df = pd.DataFrame(scaler_X.fit_transform(X_df), columns=X_df.columns) #skaalaus params and meteo else: X_df = pd.DataFrame(scaler_X.transform(X_df), columns=X_df.columns) meteo_df=X_df['meteo'] meteo_df.columns= meteo_df.columns.map(ast.literal_eval) meteo_df.columns=pd.MultiIndex.from_tuples(meteo_df.columns) meteo_df = meteo_df.swaplevel(axis=1) meteo_df.columns = meteo_df.columns.set_levels(meteo_df.columns.levels[0].astype(int), level=0) meteo_df=meteo_df.sort_index(axis=1) meteo = meteo_df.to_numpy() meteo = meteo.reshape(len(meteo), meteo_df.columns.get_level_values(0).nunique(), meteo_df.columns.get_level_values(1).nunique()) #datapoints, time steps , meteovalues param_df=X_df['parameters'] param = param_df.to_numpy() y_df=df_data['features'] y_df.columns= y_df.columns.map(ast.literal_eval) y_df.columns=pd.MultiIndex.from_tuples(y_df.columns) y_df = y_df.swaplevel(axis=1) #changing the month(tai week) level to upper y_df.columns = y_df.columns.set_levels(y_df.columns.levels[0].astype(int), level=0) y_df=y_df.sort_index(axis=1) y_n=y_df.to_numpy() if train: y_n = scaler_y.fit_transform(y_n) else: y_n = scaler_y.transform(y_n) y_n = y_n.reshape(len(y_n), y_df.columns.get_level_values(0).nunique(), y_df.columns.get_level_values(1).nunique()) if train: num_batches= meteo.shape[0] // batch_size met_batches=jnp.array_split(meteo, num_batches) param_batches=jnp.array_split(param, num_batches) y_batches=jnp.array_split(y_n, num_batches) return met_batches, param_batches, y_batches, scaler_X, scaler_y else: return meteo, param, y_n class LSTM_model(nn.Module): hidden_size_p1: Union[int, None] hidden_size_p2: Union[int, None] hidden_size_lstm: int hidden_size_comb: Union[int, None] dropout_rate: float @nn.compact def __call__(self , met, param, train:bool): cell= nn.LSTMCell(name="lstmcell", features=self.hidden_size_lstm) def body_fn(cell, carry, met): carry, y= cell(carry, met) carry = (carry[0].astype(jnp.float32), carry[1].astype(jnp.float32)) return carry, y # scan through each time step scan=nn.scan(body_fn, variable_broadcast="params", split_rngs={"params": False}, in_axes=1, out_axes=1) carry=cell.initialize_carry(jax.random.PRNGKey(0), ((met.shape)[0], self.hidden_size_lstm)) carry, y_met = scan(cell, carry, met) y_met=nn.relu(y_met) y_param=param if self.hidden_size_p1 is not None: y_param=nn.Dense(self.hidden_size_p1) (y_param) y_param=nn.relu(y_param) y_param = nn.Dropout(rate=self.dropout_rate, deterministic=not train)(y_param) if self.hidden_size_p2 is not None: y_param=nn.Dense(self.hidden_size_p2)(y_param) y_param=nn.relu(y_param) y_param = nn.Dropout(rate=self.dropout_rate, deterministic=not train)(y_param) y_param = jnp.tile(y_param[:, jnp.newaxis, :], (1, 53, 1)) y = jnp.concatenate([y_met, y_param], axis=-1) if self.hidden_size_comb is not None: y=nn.Dense(int(self.hidden_size_comb))(y) y=nn.relu(y) y = nn.Dropout(rate=self.dropout_rate, deterministic=not train)(y) y=nn.Dense(5)(y) return y def mse(params, train_state, met_batch,param_batch, y_batch, train): predicted = train_state.apply_fn({'params': params}, met_batch, param_batch, train=train, rngs={'dropout': jax.random.PRNGKey(0)}) mse_value=jnp.mean((predicted - y_batch)**2) return mse_value @jax.jit def train_step(train_state, met_batch, param_batch, y_batch): loss, grads = jax.value_and_grad(mse, allow_int=True)(train_state.params, train_state, met_batch, param_batch, y_batch, True) train_state = train_state.apply_gradients(grads=grads) return train_state, loss def create_state(met_train_batches, param_train_batches, learning_rate, model): met_shape = met_train_batches[0].shape param_shape = param_train_batches[0].shape r=jax.random.PRNGKey(0) # initialize the weights variables = model.init({'params': r}, jnp.ones(met_shape), jnp.ones(param_shape), train=False) tx=optax.adamw(learning_rate=learning_rate) # initialize the parameters of the model; weights and biases train_state = flax.training.train_state.TrainState.create(apply_fn=model.apply, params=variables['params'], tx=tx) return train_state def train(hidden_size_p1, hidden_size_p2, hidden_size_lstm, hidden_size_comb, dropout_rate, batch_size, learning_rate, epochs): model = LSTM_model(hidden_size_p1, hidden_size_p2, hidden_size_lstm, hidden_size_comb, dropout_rate) mse_fold=[] for fold in range(1, fold_num+1): train_folds = pd.concat( [pd.read_parquet(data_file_shape.format(i)) for i in range(1, fold_num + 1) if i != fold] ) met_train_batches, param_train_batches, y_train_batches, scaler_X, scaler_y=get_batches(train_folds, scaler_X=None, scaler_y=None, batch_size=batch_size, train=True) met_test, param_test, y_test = get_batches(pd.read_parquet(data_file_shape.format(fold)), scaler_X, scaler_y, batch_size, train=False) train_state = create_state(met_train_batches,param_train_batches, learning_rate, model) test_mse_list=[] train_mse_list=[] early_stop = EarlyStopping(min_delta=1e-4, patience=3) #min delta between updates to be considered an improvement & Number of steps of no improvement before stopping for epoch in range(epochs): mse_list=[] for met_b, param_b, y_b in zip(met_train_batches, param_train_batches, y_train_batches): train_state, loss=train_step(train_state, met_b ,param_b , y_b) mse_list.append(loss) mean_train_mse=jnp.mean(jnp.array(mse_list)) test_mse=mse(train_state.params, train_state, met_test, param_test, y_test, False) print(f"Fold: {fold}, Epoch: {epoch + 1}, MSE: {test_mse}, mean train MSE {mean_train_mse}") train_mse_list.append(mean_train_mse) test_mse_list.append(test_mse) early_stop=early_stop.update(mean_train_mse) if early_stop.should_stop: print(f'Met early stopping criteria, breaking at epoch {epoch+1}') break print(parameter_overview.get_parameter_overview(train_state.params)) mse_fold.append(test_mse) save_checkpoint(train_state, scaler_X, scaler_y, fold) return(mse_fold) def save_checkpoint(state, scaler_X, scaler_y, fold): ckpt_dir = 'tmp/flax_ckpt/orbax/single_save_{}'.format(fold) if os.path.exists(ckpt_dir): shutil.rmtree(ckpt_dir) ckpt= {'state': state} orbax.checkpoint.PyTreeCheckpointer().save(ckpt_dir, ckpt) #scalers with open('tmp/LSTM_scaler_X_{}.pkl'.format(fold),'wb') as f: pickle.dump(scaler_X, f) with open('tmp/LSTM_scaler_y_{}.pkl'.format(fold),'wb') as f: pickle.dump(scaler_y, f) if args.task == "train": train(hidden_size_p1=dense1_size, hidden_size_p2=dense2_size, hidden_size_lstm=LSTM_size, hidden_size_comb=dense3_size, dropout_rate=args.dropout_rate, batch_size=args.batch_size, learning_rate=args.learning_rate, epochs=args.number_epochs) elif args.task == "optimize": def objective(trial): # define hyperparameter space layer_1 = trial.suggest_categorical('layer_1', [True, False]) if layer_1 is True: hidden_size_p1 = trial.suggest_int('hidden_size_p1', 20, 400) else: hidden_size_p1 = None layer_2 = trial.suggest_categorical('layer_2', [True, False]) if layer_2 is True: hidden_size_p2 = trial.suggest_int('hidden_size_p2', 20, 400) else: hidden_size_p2 = None hidden_size_lstm = trial.suggest_int('hidden_size_lstm', 20, 400) layer_comb = trial.suggest_categorical('layer_comb', [True, False]) if layer_comb is True: hidden_size_comb = trial.suggest_int('hidden_size_comb', 20, 400) else: hidden_size_comb = None dropout_rate = trial.suggest_float('dropout_rate', 0, 0.75, step=0.01) batch_size = trial.suggest_int('batch_size', 10, 100) learning_rate = trial.suggest_categorical('learning_rate', [0.001, 0.01]) epochs = 100 print(f"Trial {trial.number}") print(f" hidden_size_p1: {hidden_size_p1}") print(f" hidden_size_p2: {hidden_size_p2}") print(f" hidden_size_lstm: {hidden_size_lstm}") print(f" hidden_size_comb: {hidden_size_comb}") print(f" dropout_rate: {dropout_rate}") print(f" batch_size: {batch_size}") print(f" learning_rate: {learning_rate}") # call the train function with the suggested hyperparameters mse_list = train(hidden_size_p1, hidden_size_p2, hidden_size_lstm, hidden_size_comb, dropout_rate, batch_size, learning_rate, epochs) mean_mse = jnp.mean(jnp.array(mse_list)) return mean_mse # create a study with TPESampler sampler = optuna.samplers.TPESampler(seed=10) study_name = 'hyperparameter_search_LSTM' storage_name = "sqlite:///{}.db".format(study_name) # create a study study = optuna.create_study( study_name=study_name, storage=storage_name, direction="minimize", sampler=sampler, load_if_exists=True ) # use if you already have created a study """study = optuna.load_study( study_name=study_name, storage=storage_name )""" study.optimize(objective, n_trials=args.number_trials) best_params = study.best_params print("Best parameters: ", best_params) with open('best_params_LSTM.txt', 'w') as f: for key, value in best_params.items(): f.write(f"{key}: {value}\n") best_value = study.best_value print("Best MSE value: ", best_value) fig=optuna.visualization.plot_optimization_history(study) fig.write_image('plots_LSTM/optimization_history.png') fig=optuna.visualization.plot_parallel_coordinate(study) fig.write_image('plots_LSTM/parallel_coordinate.png') fig = optuna.visualization.plot_optimization_history(study) fig.write_image('plots_LSTM/optimization_history.png') fig = optuna.visualization.plot_param_importances(study) fig.write_image('plots_LSTM/param_importances.png')