regr

scripts/.ipynb_checkpoints/training_reg-checkpoint.py

+import sys; sys.path.append('/home/mcristofo/DST/DST')
+import os
+from DST.config import data_path
+import pandas as pd
+import numpy as np
+import torch
+import torch.nn as nn
+import time
+import math
+import torch.utils.data as utils_data
+import torch.nn.functional as F
+import datetime
+
+from sklearn.metrics import confusion_matrix, matthews_corrcoef
+from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error
+
+np.set_printoptions(suppress=True, precision=3)
+
+torch.manual_seed(21894)
+np.random.seed(21894)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(device)
+
+BEFORE = 12
+AFTER = 12
+
+dst_data = pd.read_pickle(os.path.join(data_path,'dst.pkl'))
+dst_data['ora_round'] = dst_data.ora.apply(lambda x:int(x.split(':')[0]))
+dati_agg =  dst_data.groupby(['data','ora_round']).agg({
+                                                       'BX': np.mean,
+                                                      'BY': np.mean,
+                                                     'BZ': np.mean,
+                                                     'FLOW_SPEED': np.mean,
+                                                     'PROTON_DENSITY': np.mean,
+                                                   'TEMPERATURE': np.mean,
+                                                   'PRESSION': np.mean,   
+                                                    'ELETTRIC': np.mean,
+                                                   'y': np.mean})
+dati_agg.reset_index(inplace=True)
+dati_agg.sort_values(by = ['data','ora_round'],inplace=True)
+
+dataset = dati_agg.drop(columns = ['data','ora_round']).values
+dataset = torch.from_numpy(np.hstack([np.arange(len(dataset)).reshape([-1,1]),dataset]))
+
+last_date_train = dati_agg[dati_agg.data <= datetime.datetime(2008,12,31)].index[-1]
+len_valid_test = (len(dataset) - last_date_train)/2
+last_date_train/len(dataset), len_valid_test/len(dataset)
+
+data_in = dataset.unfold(0, BEFORE, 1).transpose(2,1)
+data_out = dataset[BEFORE:].unfold(0, AFTER, 1).transpose(2,1)
+data_in = data_in[:data_out.size(0)]
+
+data_out = data_out[:,:,-1]
+
+data_in.size(), data_out.size()
+
+where_not_nan_in = ~torch.isnan(data_in).any(2, keepdim=True).any(1, keepdim=True).reshape(-1)
+data_in = data_in[where_not_nan_in]
+data_out = data_out[where_not_nan_in]
+
+where_not_nan_out = ~torch.isnan(data_out).any(1, keepdim=True).reshape(-1)
+data_in = data_in[where_not_nan_out]
+data_out = data_out[where_not_nan_out]
+
+last_train = np.where(data_in[:,0,0] <= last_date_train)[0][-1] + 1
+
+data_in = data_in[:, :, 1:]
+
+#len_tr = int(data_in.size(0) * 0.6)
+n_channels = data_in.size(2)
+
+class MinMaxScaler():
+    """
+    Transform features by scaling each feature to a given range
+    
+    Features in the last dim
+    
+    The transformation is given by::
+        X_std = (X - X.min(axis=0)) / (X.max(axis=0) - X.min(axis=0))
+        X_scaled = X_std * (max - min) + min
+    where min, max = feature_range.    
+    """
+    
+    def __init__(self, feature_range=(0,1)):
+        self.feature_range = feature_range
+    
+    def fit(self, X):
+        X_size = X.size()
+        X = X.reshape(-1, X_size[-1])
+        data_min = X.min(axis=0).values
+        data_max = X.max(axis=0).values
+        
+        data_range = data_max - data_min
+        self.scale_ = ((self.feature_range[1] - self.feature_range[0]) / data_range)
+        self.min_ = self.feature_range[0] - data_min * self.scale_
+        self.data_min_ = data_min
+        self.data_max_ = data_max
+        self.data_range_ = data_range
+        
+        X = X.reshape(X_size)
+        return self
+    
+    def transform(self, X):
+        X *= self.scale_
+        X += self.min_
+        return X 
+    
+    def inverse_transform(self, X):
+
+        X -= self.min_
+        X /= self.scale_
+        return X        
+        
+        
+mmScaler = MinMaxScaler((0.1, .9))
+mmScaler.fit(data_in[:last_train])
+data_in_scaled = data_in.clone()
+data_in_scaled = mmScaler.transform(data_in_scaled)
+
+mm_scaler_out = MinMaxScaler((0.1, .9))
+mm_scaler_out.fit(data_in[:last_train, :, -1].reshape(-1, data_in.size(1), 1))
+data_out_scaled = data_out.clone()
+data_out_scaled = mm_scaler_out.transform(data_out_scaled)
+
+
+dst_levels = [-20,-50,-100]
+data_out_c = data_out.clone()
+data_out_c[np.where(data_out_c >= dst_levels[0])] = 0
+data_out_c[np.where((data_out_c < dst_levels[0]) & (data_out_c >= dst_levels[1]))] = 1
+data_out_c[np.where((data_out_c < dst_levels[1]) & (data_out_c >= dst_levels[2]))] = 2
+data_out_c[np.where((data_out_c < dst_levels[2]))] = 3
+
+class Dataset(utils_data.Dataset):
+    
+    def __init__(self, dataset_in, dataset_out):
+        self.dataset_in = dataset_in        
+        self.dataset_out = dataset_out 
+        
+    def __len__(self):
+        return self.dataset_in.size(0)
+
+    def __getitem__(self, idx):
+
+        din_src = self.dataset_in[idx]
+        dout = self.dataset_out[idx]
+
+        return din_src, dout
+    
+ixs_valid_test = np.arange(int(len_valid_test)) + last_train
+np.random.shuffle(ixs_valid_test)
+ixs_valid = ixs_valid_test[::2]
+ixs_test = ixs_valid_test[1::2]
+
+ixs_tr1 = np.where((mm_scaler_out.inverse_transform(data_out_scaled[:last_train].clone()).min(axis=1)[0]).numpy()<-20)[0]
+ixs_tr2 = np.where((mm_scaler_out.inverse_transform(data_out_scaled[:last_train].clone()).min(axis=1)[0]).numpy()>=-20)[0]
+
+BATCH_SIZE=256
+
+np.random.shuffle(ixs_tr2)
+ixs_tr2a = ixs_tr2[:10000]
+ixs_tr = list(ixs_tr1) + list(ixs_tr2a)
+dst_min = data_out[ixs_tr].min(axis=1).values.flatten()
+
+bins = [dst_min.min() - 10] + list(np.arange(-300, dst_min.max() + 10, 10))
+h, b = np.histogram(dst_min, bins=bins)
+if len(np.argwhere(h == 0)) > 0:
+    bins = np.delete(bins, np.argwhere(h == 0)[0] + 1)
+    h, b = np.histogram(dst_min, bins=bins)
+w = h.max()/h
+
+def fix_weight(dst_v):
+    pos = np.argwhere(np.abs(b - dst_v) == np.abs((b - dst_v)).min())[0,0]
+    if dst_v - b[pos] < 0:
+        pos = pos-1
+    return np.sqrt(w[pos]/h.max())
+#     return w[pos]/h.max()
+
+fix_weight_v = np.vectorize(fix_weight)    
+weights = fix_weight_v(dst_min)
+
+sampler = torch.utils.data.sampler.WeightedRandomSampler(weights, num_samples= len(dst_min))
+
+dataset_tr = Dataset(data_in_scaled[ixs_tr], data_out_scaled[ixs_tr])
+data_loader_tr = utils_data.DataLoader(dataset_tr, batch_size=BATCH_SIZE, num_workers = 4, shuffle=False, sampler = sampler)
+
+class DSTnet(nn.Module):
+    def __init__(self, nvars, nhidden_i, nhidden_o, n_out_i, before, after):
+        super().__init__()
+        self.nvars = nvars
+        self.nhidden_i = nhidden_i
+        self.nhidden_o = nhidden_o
+        self.before = before
+        self.after = after
+        self.n_out_i = n_out_i
+
+        self.lstm = nn.LSTM(self.nvars, self.n_out_i, self.nhidden_i, batch_first=True)
+        
+        self.first_merged_layer = self.n_out_i * self.before
+        self.bn1 = nn.BatchNorm1d(num_features=self.first_merged_layer)
+#         self.bn1 = nn.LayerNorm(self.first_merged_layer)
+        self.linear_o_1 = nn.Linear(self.first_merged_layer, self.nhidden_o)
+        self.ln1 = nn.LayerNorm(self.nhidden_o )
+        self.linear_o_2 = nn.Linear(self.nhidden_o, self.nhidden_o)
+        self.linear_o_3 = nn.Linear(self.nhidden_o, self.nhidden_o // 2)
+#         self.bn2 = nn.BatchNorm1d(num_features=self.nhidden_o)
+        self.linear_o_4 = nn.Linear(self.nhidden_o // 2, self.after)
+    
+    def init_hidden(self, batch_size):
+        hidden = torch.randn(self.nhidden_i, batch_size, self.n_out_i).to(device)
+        cell =  torch.randn(self.nhidden_i, batch_size, self.n_out_i).to(device)
+        return (hidden, cell)
+    
+        
+    def forward(self, x0):
+        
+        self.hidden = self.init_hidden(x0.size(0))
+        x = self.lstm(x0, self.hidden)[0].reshape(x0.shape[0], -1)
+        x = self.bn1(x)
+#         x = F.relu(x)
+        x = F.relu(self.linear_o_1(x))
+#         x = self.ln1(x)
+        x = F.dropout(x, 0.2, training=self.training)
+        x = F.relu(self.linear_o_2(x))
+        x = F.dropout(x, 0.2, training=self.training)
+        x = F.relu(self.linear_o_3(x))
+        x = F.dropout(x, 0.2, training=self.training)
+        x = self.linear_o_4(x)
+        return x
+    
+loss_f = nn.L1Loss()
+loss_mse = nn.MSELoss()
+
+
+nhidden_i = 1
+nhidden_o = 64
+n_out_i = 16
+before = BEFORE
+nvars = data_in_scaled.shape[-1]
+
+dst_net = DSTnet(nvars, nhidden_i, nhidden_o, n_out_i, before, AFTER).to(device)
+print(dst_net)
+
+num_epochs = 10000
+lr = 1e-5
+optimizer = torch.optim.Adam(dst_net.parameters(), lr=lr, weight_decay=1e-5)
+
+history_tr = np.zeros((num_epochs, 2))
+history_valid = np.zeros((num_epochs, 2))
+history_ts = np.zeros((num_epochs, 2))
+
+for epoch in range(num_epochs):
+    start_time = time.time()
+
+    for i, batch in enumerate(data_loader_tr):
+        
+        x = (batch[0] + (( 1- 2 * torch.rand(batch[0].size())) * batch[0] * 0.001)).float().to(device) #+ delta_batch0
+        y = (batch[1] + (( 1- 2 * torch.rand(batch[1].size())) * batch[1] * 0.001)).float().to(device) #+ delta_batch1
+        optimizer.zero_grad()
+        dst_net.train()
+        outputs = dst_net(x)
+        loss = loss_f(outputs, y) #+ torch.sqrt(loss_mse(outputs, y) + 0.0000001)# * (1 + torch.randn(y.shape).to(device) * 0.01))
+
+        loss.backward()
+
+        optimizer.step()
+
+    dst_net.eval()
+    data_out_scaled_loss = mm_scaler_out.inverse_transform(data_out_scaled.clone())
+    outputs = dst_net(data_in_scaled[:last_train].to(device).float())
+    loss_tr = np.sqrt(loss_mse(mm_scaler_out.inverse_transform(outputs.cpu().clone()).to(device), data_out_scaled_loss[:last_train].to(device).float()).item())
+    loss_mae_tr = loss_f(mm_scaler_out.inverse_transform(outputs.cpu().clone()).to(device), data_out_scaled_loss[:last_train].to(device).float()).item()
+
+    outputs = dst_net(data_in_scaled[ixs_valid].to(device).float())
+    loss_valid = np.sqrt(loss_mse(mm_scaler_out.inverse_transform(outputs.cpu().clone()).to(device), data_out_scaled_loss[ixs_valid].to(device).float()).item())
+    loss_mae_valid = loss_f(mm_scaler_out.inverse_transform(outputs.cpu().clone()).to(device), data_out_scaled_loss[ixs_valid].to(device).float()).item()
+
+    outputs = dst_net(data_in_scaled[ixs_test].to(device).float())
+    loss_ts = np.sqrt(loss_mse(mm_scaler_out.inverse_transform(outputs.cpu().clone()).to(device), data_out_scaled_loss[ixs_test].to(device).float()).item())
+    loss_mae_ts = loss_f(mm_scaler_out.inverse_transform(outputs.cpu().clone()).to(device), data_out_scaled_loss[ixs_test].to(device).float()).item()
+
+
+    history_tr[epoch] = [loss_tr, loss_mae_tr]
+    history_valid[epoch] = [loss_valid, loss_mae_valid]
+    history_ts[epoch] = [loss_ts, loss_mae_ts]
+    epoch_time = time.time() - start_time
+
+    if (epoch % 1 == 0):
+        print('Epoch %d time = %.2f, tr_rmse = %0.5f, valid_rmse = %.5f, ts_rmse = %.5f, tr_mae = %0.5f, valid_mae = %.5f, ts_mae = %.5f' % 
+              (epoch, epoch_time, loss_tr, loss_valid, loss_ts, loss_mae_tr, loss_mae_valid, loss_mae_ts))
+
+    if (epoch % 10 == 0):
+
+        out_r = dst_net(data_in_scaled[last_train:].to(device).float())
+
+        tp = 11
+
+        dst_levels = [-20,-50,-100]
+
+        truth = data_out[last_train:].cpu().detach().numpy().copy()
+        out = mm_scaler_out.inverse_transform(out_r.cpu().clone()).detach().numpy()
+
+        for ll in range(12):
+            print(ll, np.sqrt(mean_squared_error(truth[:,ll], out[:,ll])))
+
+        truth[np.where(truth >= dst_levels[0])] = 0
+        truth[np.where((truth < dst_levels[0]) & (truth >= dst_levels[1]))] = 1
+        truth[np.where((truth < dst_levels[1]) & (truth >= dst_levels[2]))] = 2
+        truth[np.where((truth < dst_levels[2]))] = 3
+
+        out[np.where(out >= dst_levels[0])] = 0
+        out[np.where((out < dst_levels[0]) & (out >= dst_levels[1]))] = 1
+        out[np.where((out < dst_levels[1]) & (out >= dst_levels[2]))] = 2
+        out[np.where((out < dst_levels[2]))] = 3
+
+        print(confusion_matrix(truth[:,11], out[:,11])/confusion_matrix(truth[:,11], out[:,11]).sum(axis=1)[:, None])
+
+        np.random.shuffle(ixs_tr2)
+        ixs_tr2a = ixs_tr2[:10000]
+        ixs_tr = list(ixs_tr1) + list(ixs_tr2a)
+        dst_min = data_out[ixs_tr].min(axis=1).values.flatten()
+
+        bins = [dst_min.min() - 10] + list(np.arange(-300, dst_min.max() + 10, 10))
+        h, b = np.histogram(dst_min, bins=bins)
+        if len(np.argwhere(h == 0)) > 0:
+            bins = np.delete(bins, np.argwhere(h == 0)[0] + 1)
+            h, b = np.histogram(dst_min, bins=bins)
+        w = h.max()/h
+
+        fix_weight_v = np.vectorize(fix_weight)    
+        weights = fix_weight_v(dst_min)
+
+        sampler = torch.utils.data.sampler.WeightedRandomSampler(weights, num_samples= len(dst_min))
+        dataset_tr = Dataset(data_in_scaled[ixs_tr], data_out_scaled[ixs_tr])
+        data_loader_tr = utils_data.DataLoader(dataset_tr, batch_size=BATCH_SIZE, num_workers = 4, shuffle=False, sampler = sampler)        
+
+torch.save(dst_net.state_dict(), 'models/dst_reg_64_16.pth')    
+
+np.savetxt('hist/history_tr_rmse_mae_reg_64_16.txt', history_tr)
+np.savetxt('hist/history_valid_rmse_mae_reg_64_16.txt', history_valid)
+np.savetxt('hist/history_ts_rmse_mae_reg_64_16.txt', history_ts)

scripts/training_reg.py

+import sys; sys.path.append('/home/mcristofo/DST/DST')
+import os
+from DST.config import data_path
+import pandas as pd
+import numpy as np
+import torch
+import torch.nn as nn
+import time
+import math
+import torch.utils.data as utils_data
+import torch.nn.functional as F
+import datetime
+
+from sklearn.metrics import confusion_matrix, matthews_corrcoef
+from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error
+
+np.set_printoptions(suppress=True, precision=3)
+
+torch.manual_seed(21894)
+np.random.seed(21894)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(device)
+
+BEFORE = 12
+AFTER = 12
+
+dst_data = pd.read_pickle(os.path.join(data_path,'dst.pkl'))
+dst_data['ora_round'] = dst_data.ora.apply(lambda x:int(x.split(':')[0]))
+dati_agg =  dst_data.groupby(['data','ora_round']).agg({
+                                                       'BX': np.mean,
+                                                      'BY': np.mean,
+                                                     'BZ': np.mean,
+                                                     'FLOW_SPEED': np.mean,
+                                                     'PROTON_DENSITY': np.mean,
+                                                   'TEMPERATURE': np.mean,
+                                                   'PRESSION': np.mean,   
+                                                    'ELETTRIC': np.mean,
+                                                   'y': np.mean})
+dati_agg.reset_index(inplace=True)
+dati_agg.sort_values(by = ['data','ora_round'],inplace=True)
+
+dataset = dati_agg.drop(columns = ['data','ora_round']).values
+dataset = torch.from_numpy(np.hstack([np.arange(len(dataset)).reshape([-1,1]),dataset]))
+
+last_date_train = dati_agg[dati_agg.data <= datetime.datetime(2008,12,31)].index[-1]
+len_valid_test = (len(dataset) - last_date_train)/2
+last_date_train/len(dataset), len_valid_test/len(dataset)
+
+data_in = dataset.unfold(0, BEFORE, 1).transpose(2,1)
+data_out = dataset[BEFORE:].unfold(0, AFTER, 1).transpose(2,1)
+data_in = data_in[:data_out.size(0)]
+
+data_out = data_out[:,:,-1]
+
+data_in.size(), data_out.size()
+
+where_not_nan_in = ~torch.isnan(data_in).any(2, keepdim=True).any(1, keepdim=True).reshape(-1)
+data_in = data_in[where_not_nan_in]
+data_out = data_out[where_not_nan_in]
+
+where_not_nan_out = ~torch.isnan(data_out).any(1, keepdim=True).reshape(-1)
+data_in = data_in[where_not_nan_out]
+data_out = data_out[where_not_nan_out]
+
+last_train = np.where(data_in[:,0,0] <= last_date_train)[0][-1] + 1
+
+data_in = data_in[:, :, 1:]
+
+#len_tr = int(data_in.size(0) * 0.6)
+n_channels = data_in.size(2)
+
+class MinMaxScaler():
+    """
+    Transform features by scaling each feature to a given range
+    
+    Features in the last dim
+    
+    The transformation is given by::
+        X_std = (X - X.min(axis=0)) / (X.max(axis=0) - X.min(axis=0))
+        X_scaled = X_std * (max - min) + min
+    where min, max = feature_range.    
+    """
+    
+    def __init__(self, feature_range=(0,1)):
+        self.feature_range = feature_range
+    
+    def fit(self, X):
+        X_size = X.size()
+        X = X.reshape(-1, X_size[-1])
+        data_min = X.min(axis=0).values
+        data_max = X.max(axis=0).values
+        
+        data_range = data_max - data_min
+        self.scale_ = ((self.feature_range[1] - self.feature_range[0]) / data_range)
+        self.min_ = self.feature_range[0] - data_min * self.scale_
+        self.data_min_ = data_min
+        self.data_max_ = data_max
+        self.data_range_ = data_range
+        
+        X = X.reshape(X_size)
+        return self
+    
+    def transform(self, X):
+        X *= self.scale_
+        X += self.min_
+        return X 
+    
+    def inverse_transform(self, X):
+
+        X -= self.min_
+        X /= self.scale_
+        return X        
+        
+        
+mmScaler = MinMaxScaler((0.1, .9))
+mmScaler.fit(data_in[:last_train])
+data_in_scaled = data_in.clone()
+data_in_scaled = mmScaler.transform(data_in_scaled)
+
+mm_scaler_out = MinMaxScaler((0.1, .9))
+mm_scaler_out.fit(data_in[:last_train, :, -1].reshape(-1, data_in.size(1), 1))
+data_out_scaled = data_out.clone()
+data_out_scaled = mm_scaler_out.transform(data_out_scaled)
+
+
+dst_levels = [-20,-50,-100]
+data_out_c = data_out.clone()
+data_out_c[np.where(data_out_c >= dst_levels[0])] = 0
+data_out_c[np.where((data_out_c < dst_levels[0]) & (data_out_c >= dst_levels[1]))] = 1
+data_out_c[np.where((data_out_c < dst_levels[1]) & (data_out_c >= dst_levels[2]))] = 2
+data_out_c[np.where((data_out_c < dst_levels[2]))] = 3
+
+class Dataset(utils_data.Dataset):
+    
+    def __init__(self, dataset_in, dataset_out):
+        self.dataset_in = dataset_in        
+        self.dataset_out = dataset_out 
+        
+    def __len__(self):
+        return self.dataset_in.size(0)
+
+    def __getitem__(self, idx):
+
+        din_src = self.dataset_in[idx]
+        dout = self.dataset_out[idx]
+
+        return din_src, dout
+    
+ixs_valid_test = np.arange(int(len_valid_test)) + last_train
+np.random.shuffle(ixs_valid_test)
+ixs_valid = ixs_valid_test[::2]
+ixs_test = ixs_valid_test[1::2]
+
+ixs_tr1 = np.where((mm_scaler_out.inverse_transform(data_out_scaled[:last_train].clone()).min(axis=1)[0]).numpy()<-20)[0]
+ixs_tr2 = np.where((mm_scaler_out.inverse_transform(data_out_scaled[:last_train].clone()).min(axis=1)[0]).numpy()>=-20)[0]
+
+BATCH_SIZE=256
+
+np.random.shuffle(ixs_tr2)
+ixs_tr2a = ixs_tr2[:10000]
+ixs_tr = list(ixs_tr1) + list(ixs_tr2a)
+dst_min = data_out[ixs_tr].min(axis=1).values.flatten()
+
+bins = [dst_min.min() - 10] + list(np.arange(-300, dst_min.max() + 10, 10))
+h, b = np.histogram(dst_min, bins=bins)
+if len(np.argwhere(h == 0)) > 0:
+    bins = np.delete(bins, np.argwhere(h == 0)[0] + 1)
+    h, b = np.histogram(dst_min, bins=bins)
+w = h.max()/h
+
+def fix_weight(dst_v):
+    pos = np.argwhere(np.abs(b - dst_v) == np.abs((b - dst_v)).min())[0,0]
+    if dst_v - b[pos] < 0:
+        pos = pos-1
+    return np.sqrt(w[pos]/h.max())
+#     return w[pos]/h.max()
+
+fix_weight_v = np.vectorize(fix_weight)    
+weights = fix_weight_v(dst_min)
+
+sampler = torch.utils.data.sampler.WeightedRandomSampler(weights, num_samples= len(dst_min))
+
+dataset_tr = Dataset(data_in_scaled[ixs_tr], data_out_scaled[ixs_tr])
+data_loader_tr = utils_data.DataLoader(dataset_tr, batch_size=BATCH_SIZE, num_workers = 4, shuffle=False, sampler = sampler)
+
+class DSTnet(nn.Module):
+    def __init__(self, nvars, nhidden_i, nhidden_o, n_out_i, before, after):
+        super().__init__()
+        self.nvars = nvars
+        self.nhidden_i = nhidden_i
+        self.nhidden_o = nhidden_o
+        self.before = before
+        self.after = after
+        self.n_out_i = n_out_i
+
+        self.lstm = nn.LSTM(self.nvars, self.n_out_i, self.nhidden_i, batch_first=True)
+        
+        self.first_merged_layer = self.n_out_i * self.before
+        self.bn1 = nn.BatchNorm1d(num_features=self.first_merged_layer)
+#         self.bn1 = nn.LayerNorm(self.first_merged_layer)
+        self.linear_o_1 = nn.Linear(self.first_merged_layer, self.nhidden_o)
+        self.ln1 = nn.LayerNorm(self.nhidden_o )
+        self.linear_o_2 = nn.Linear(self.nhidden_o, self.nhidden_o)
+        self.linear_o_3 = nn.Linear(self.nhidden_o, self.nhidden_o // 2)
+#         self.bn2 = nn.BatchNorm1d(num_features=self.nhidden_o)
+        self.linear_o_4 = nn.Linear(self.nhidden_o // 2, self.after)
+    
+    def init_hidden(self, batch_size):
+        hidden = torch.randn(self.nhidden_i, batch_size, self.n_out_i).to(device)
+        cell =  torch.randn(self.nhidden_i, batch_size, self.n_out_i).to(device)
+        return (hidden, cell)
+    
+        
+    def forward(self, x0):
+        
+        self.hidden = self.init_hidden(x0.size(0))
+        x = self.lstm(x0, self.hidden)[0].reshape(x0.shape[0], -1)
+        x = self.bn1(x)
+#         x = F.relu(x)
+        x = F.relu(self.linear_o_1(x))
+#         x = self.ln1(x)
+        x = F.dropout(x, 0.2, training=self.training)
+        x = F.relu(self.linear_o_2(x))