added skip connection to FFNN

pyforecaster/forecasting_models/neural_models/INN.py

@@ -80,7 +80,7 @@ class EndToEndCausalInvertibleModule(nn.Module):
     activation: callable = nn.relu
     def setup(self):
         self.embedder = CausalInvertibleModule(num_layers=self.num_embedding_layers, features=self.features_embedding)
-        self.predictor = FeedForwardModule(n_layers=np.hstack([(np.ones(self.num_prediction_layers-1)*self.features_prediction).astype(int), self.n_out]), split_heads=True)
+        self.predictor = FeedForwardModule(n_layers=np.hstack([(np.ones(self.num_prediction_layers-1)*self.features_prediction).astype(int), self.n_out]), split_heads=False)
         self.invert_fun = jax.jit(partial(self.inverter, embedder=self.embedder))
     def __call__(self, x):
         x = x.copy()

pyforecaster/forecasting_models/neural_models/base_nn.py

@@ -12,6 +12,7 @@
 import pandas as pd
 from flax import linen as nn
 from functools import partial
+
 from jax import random, value_and_grad, vmap
 from jax.scipy.special import erfinv
 from scipy.special import erfinv
@@ -97,18 +98,20 @@ def __call__(self, x):
             layers = layers.astype(int)
         else:
             layers = self.n_layers
+        y = nn.Dense(features=layers[-1], name='dense_-1')(x)
         for i, n in enumerate(layers):
-
             if i < len(layers)-1:
                 x = nn.Dense(features=n, name='dense_{}'.format(i))(x)
                 x = nn.relu(x)
+                x = nn.LayerNorm()(x)
             else:
                 if self.split_heads:
                     n_out = self.n_out if self.n_out is not None else layers[-1]
                     # split into n_out heads to predict the output independently
                     subnets = [nn.relu(nn.Dense(features=layers[np.maximum(-2, -len(layers))], name='subnet_in_{}'.format(k))(x)) for k in range(n_out)]
                     out = [nn.Dense(features=1, name='subnet_out_{}'.format(k))(subnets[k]) for k in range(n_out)]
                     x = jnp.hstack(out)
+
                     """
                     # Combine the outputs in a single dense layer
                     n_last = layers[np.maximum(-2, -len(layers))]
@@ -123,7 +126,7 @@ def __call__(self, x):
                 else:
                     x = nn.Dense(features=n, name='dense_{}'.format(i))(x)
 
-        return x
+        return x + y
 
 class NN(ScenarioGenerator):
     input_scaler: StandardScaler = None
@@ -404,14 +407,15 @@ def get_normalized_inputs(self, inputs, target=None):
 
 class FFNN(NN):
     def __init__(self, n_out=None, q_vect=None, n_epochs=10, val_ratio=None, nodes_at_step=None, learning_rate=1e-3,
-                       scengen_dict={}, batch_size=None, **model_kwargs):
+                       scengen_dict={}, batch_size=None, split_heads=False, **model_kwargs):
+        self.split_heads = split_heads
         super().__init__(n_out=n_out, q_vect=q_vect, n_epochs=n_epochs, val_ratio=val_ratio, nodes_at_step=nodes_at_step, learning_rate=learning_rate,
                  nn_module=FeedForwardModule, scengen_dict=scengen_dict, batch_size=batch_size,  **model_kwargs)
 
     def set_arch(self):
         self.optimizer = optax.adamw(learning_rate=self.learning_rate)
         self.model = FeedForwardModule(n_layers=self.n_layers, n_neurons=self.n_hidden_x,
-                              n_out=self.n_out, split_heads=True)
+                              n_out=self.n_out, split_heads=self.split_heads)
         self.predict_batch = vmap(jitting_wrapper(predict_batch, self.model), in_axes=(None, 0))
         self.loss_fn = jitting_wrapper(probabilistic_loss_fn, self.predict_batch) if self.probabilistic else (
             jitting_wrapper(loss_fn, self.predict_batch))

tests/test_nns.py

@@ -437,8 +437,10 @@ def test_invertible_causal_nn(self):
         # if not there, create directory savepath_tr_plots
         if not exists(savepath_tr_plots):
             makedirs(savepath_tr_plots)
-        n_in = 145
+
         n_out = 10
+        n_in = e_tr.shape[1]  - n_out
+
         m_lin = LinearForecaster().fit(e_tr.iloc[:, :n_in], e_tr.iloc[:, -n_out:])
         y_hat_lin = m_lin.predict(e_te.iloc[:, :n_in])
         """
@@ -498,9 +500,21 @@ def test_invertible_causal_nn(self):
             ax.plot(y_invert.iloc[i, 144:].values, linestyle='--')
             plt.pause(1e-6)
         """
-        m = FFNN(n_hidden_x=50, n_layers=1, learning_rate=1e-3, batch_size=100, load_path=None, n_out=n_out, rel_tol=-1, stopping_rounds=20,n_epochs=1).fit(e_tr.iloc[:, :n_in], e_tr.iloc[:, -n_out:])
+        m = FFNN(n_hidden_x=20, n_layers=2, learning_rate=1e-3, batch_size=500,
+                 load_path=None, n_out=n_out, rel_tol=-1, stopping_rounds=10, n_epochs=10).fit(e_tr.iloc[:, :n_in], e_tr.iloc[:, -n_out:])
         y_hat = m.predict(e_te.iloc[:, :n_in])
 
+        from pyforecaster.metrics import summary_scores, nmae
+
+        maes_nn = summary_scores(y_hat, e_te.iloc[:, -n_out:], metrics=[nmae],idxs=pd.DataFrame(e_te.index.hour, index=e_te.index))
+        maes_lin = summary_scores(y_hat_lin, e_te.iloc[:, -n_out:], metrics=[nmae],idxs=pd.DataFrame(e_te.index.hour, index=e_te.index))
+
+        fig, ax = plt.subplots(1, 1)
+        ax.plot(maes_lin['nmae'].mean(axis=0))
+        ax.plot(maes_nn['nmae'].mean(axis=0))
+
+
+
         m = CausalInvertibleNN(learning_rate=1e-2, batch_size=300, load_path=None, n_in=n_in,
                                n_layers=2, normalize_target=False, n_epochs=5, stopping_rounds=30, rel_tol=-1,
                                end_to_end='full', n_hidden_y=300, n_prediction_layers=3, n_out=n_out,names_exogenous=['all_lag_000']).fit(e_tr.iloc[:, :n_in], e_tr.iloc[:, -n_out:])