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) # fit and scale training inputs else: X_df = pd.DataFrame(scaler_X.transform(X_df), columns=X_df.columns) # scale testing inputs with scale fitted for training data # data processing 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()) 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) 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) # fit and scale train targets else: y_n = scaler_y.transform(y_n) # scale test targets with scale fitted for train targets y_n = y_n.reshape(len(y_n), y_df.columns.get_level_values(0).nunique(), y_df.columns.get_level_values(1).nunique()) # divide the data in to batches for training 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) # LSTM cell 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 # first FNN layer 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) # second FNN layer 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) # concatenate parameters and meteorological input y_param = jnp.tile(y_param[:, jnp.newaxis, :], (1, 53, 1)) y = jnp.concatenate([y_met, y_param], axis=-1) # FNN layer to refine the LSTM output 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) # final FNN layer to produce the final output 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): # define model with proposed hyperparameters model = LSTM_model(hidden_size_p1, hidden_size_p2, hidden_size_lstm, hidden_size_comb, dropout_rate) mse_fold=[] # loop through folds to perform cross validation, where fold performs as test data and other folds as training data 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] ) # standardizing and train data to batches 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 improvement between updates to be considered an improvement (min_delta) & number of steps of no improvement before stopping (patience) # train model for the defined number of epochs 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 and want to continue from that """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) # visualize optimization 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_param_importances(study) fig.write_image('plots_LSTM/param_importances.png')