dsipts 1.1.10__tar.gz → 1.1.12__tar.gz

{dsipts-1.1.10 → dsipts-1.1.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dsipts
-Version: 1.1.10
+Version: 1.1.12
 Summary: Unified library for timeseries modelling
 Author-email: Andrea Gobbi <agobbi@fbk.eu>
 Project-URL: Homepage, https://github.com/DSIP-FBK/DSIPTS

{dsipts-1.1.10 → dsipts-1.1.12}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "dsipts"
-version = "1.1.10"
+version = "1.1.12"
 description = "Unified library for timeseries modelling"
 readme = "README.md"
 requires-python = "==3.11.13"

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/data_structure/data_structure.py

@@ -35,7 +35,18 @@ from .modifiers import *
 from aim.pytorch_lightning import AimLogger
 import time
 
-
+class DummyScaler():
+    def __init__(self):
+        pass
+    def fit(self,x):
+        pass
+    def transform(self,x):
+        return x
+    def inverse_transform(self,x):
+        return x
+    def fit_transform(self,x):
+        return x
+    
 
 pd.options.mode.chained_assignment = None 
 log = logging.getLogger(__name__)
@@ -210,20 +221,23 @@ class TimeSeries():
         self.future_variables = []
         self.target_variables = ['signal']
         self.num_var = list(set(self.past_variables).union(set(self.future_variables)).union(set(self.target_variables)))
-        
+        self.num_var = list(np.sort(self.num_var))
         
     def enrich(self,dataset,columns):
-        if columns =='hour':
-            dataset[columns] = dataset.time.dt.hour
-        elif columns=='dow':
-            dataset[columns] = dataset.time.dt.weekday
-        elif columns=='month':
-            dataset[columns] = dataset.time.dt.month
-        elif columns=='minute':
-            dataset[columns] = dataset.time.dt.minute
-        else:
-            if columns  not in dataset.columns:
-                beauty_string(f'I can not automatically enrich column {columns}. Please contact the developers or add it manually to your dataset.','section',True)
+        try:
+            if columns =='hour':
+                dataset[columns] = dataset.time.dt.hour
+            elif columns=='dow':
+                dataset[columns] = dataset.time.dt.weekday
+            elif columns=='month':
+                dataset[columns] = dataset.time.dt.month
+            elif columns=='minute':
+                dataset[columns] = dataset.time.dt.minute
+            else:
+                if columns  not in dataset.columns:
+                    beauty_string(f'I can not automatically enrich column {columns}. Please contact the developers or add it manually to your dataset.','section',True)
+        except:
+            beauty_string(f'I can not automatically enrich column {columns}. Probably not a temporal index.','section',True)
 
     def load_signal(self,data:pd.DataFrame,
                     enrich_cat:List[str] = [],
@@ -300,7 +314,7 @@ class TimeSeries():
             if check_past:
                 beauty_string('I will update past column adding all target columns, if you want to avoid this beahviour please use check_pass as false','info',self.verbose)
                 past_variables = list(set(past_variables).union(set(target_variables)))
-
+                past_variables = list(np.sort(past_variables))
         self.cat_past_var = cat_past_var
         self.cat_fut_var = cat_fut_var
 
@@ -321,14 +335,18 @@ class TimeSeries():
                 beauty_string('Categorical {c} already present, it will be added to categorical variable but not call the enriching function','info',self.verbose) 
             else:
                 self.enrich(dataset,c)
+        self.cat_past_var = list(np.sort(self.cat_past_var))
+        self.cat_fut_var = list(np.sort(self.cat_fut_var))
+        
         self.cat_var = list(set(self.cat_past_var+self.cat_fut_var)) ## all categorical data
-
+        self.cat_var = list(np.sort(self.cat_var))
         self.dataset = dataset
         self.past_variables = past_variables
         self.future_variables = future_variables
         self.target_variables = target_variables
         self.out_vars = len(target_variables)
         self.num_var = list(set(self.past_variables).union(set(self.future_variables)).union(set(self.target_variables)))
+        self.num_var = list(np.sort(self.num_var))
         if silly_model:
             beauty_string('YOU ARE TRAINING A SILLY MODEL WITH THE TARGETS IN THE INPUTS','section',self.verbose) 
             self.future_variables+=self.target_variables
@@ -665,7 +683,8 @@ class TimeSeries():
         #self.model.apply(weight_init_zeros)
 
         self.config = config
-        
+
+
         beauty_string('Setting the model','block',self.verbose)
         beauty_string(model,'',self.verbose)
         
@@ -790,8 +809,17 @@ class TimeSeries():
             weight_exists = False
             beauty_string('I can not load a previous model','section',self.verbose)
 
+        self.model.to(torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
+        if self.model.can_be_compiled():
+            try:
+                self.model = torch.compile(self.model)
+                beauty_string('Model COMPILED','block',self.verbose)
+
+            except:
+                beauty_string('Can not compile the model','block',self.verbose)
+        else:
+            beauty_string('Model can not still be compiled, be patient','block',self.verbose)
 
-        
         
         if OLD_PL:
             trainer = pl.Trainer(default_root_dir=dirpath,
@@ -873,10 +901,19 @@ class TimeSeries():
             self.losses = pd.DataFrame()
 
         try:
+
             if OLD_PL:
-                self.model = self.model.load_from_checkpoint(self.checkpoint_file_last)
+                if isinstance(self.model, torch._dynamo.eval_frame.OptimizedModule):
+                    self.model = self.model._orig_mod 
+                    self.model.load_from_checkpoint(self.checkpoint_file_last)
+                else:
+                    self.model = self.model.load_from_checkpoint(self.checkpoint_file_last)
             else:
-                self.model = self.model.__class__.load_from_checkpoint(self.checkpoint_file_last)
+                if isinstance(self.model, torch._dynamo.eval_frame.OptimizedModule):
+                    mm = self.model._orig_mod 
+                    self.model = mm.__class__.load_from_checkpoint(self.checkpoint_file_last)
+                else:
+                    self.model = self.model.__class__.load_from_checkpoint(self.checkpoint_file_last)
 
         except Exception as _:
             beauty_string(f'There is a problem loading the weights on file MAYBE CHANGED HOW WEIGHTS ARE LOADED {self.checkpoint_file_last}','section',self.verbose)
@@ -1164,6 +1201,6 @@ class TimeSeries():
                 self.model = self.model.load_from_checkpoint(tmp_path,verbose=self.verbose,)
             else:
                 self.model = self.model.__class__.load_from_checkpoint(tmp_path,verbose=self.verbose,)
-
+            self.model.to(torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
         except Exception as e:
             beauty_string(f'There is a problem loading the weights on file {tmp_path} {e}','section',self.verbose)

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/Autoformer.py

@@ -148,7 +148,8 @@ class Autoformer(Base):
             projection=nn.Linear(d_model, self.out_channels*self.mul, bias=True)
         )
         self.projection = nn.Linear(self.past_channels,self.out_channels*self.mul )
-
+    def can_be_compiled(self):
+        return True
     def forward(self, batch):
         
 

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/CrossFormer.py

@@ -114,7 +114,8 @@ class CrossFormer(Base):
         
                                          
         
-     
+    def can_be_compiled(self):
+        return True
         
     def forward(self, batch):
 

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/D3VAE.py

@@ -98,7 +98,8 @@ class D3VAE(Base):
         self.gamma = 0.01
         self.lambda1 = 1.0
 
-        
+    def can_be_compiled(self):
+        return False   
     def configure_optimizers(self):
         """
         Each model has optim_config and scheduler_config

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/Diffusion.py

@@ -425,6 +425,9 @@ class Diffusion(Base):
         loss = self.compute_loss(batch,out)
         return loss
 
+    def can_be_compiled(self):
+        return False
+
     # function to concat embedded categorical variables
     def cat_categorical_vars(self, batch:dict):
         """Extracting categorical context about past and future

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/DilatedConv.py

@@ -234,7 +234,8 @@ class DilatedConv(Base):
         self.return_additional_loss = True
         
 
-        
+    def can_be_compiled(self):
+        return True
 
     def forward(self, batch):
         """It is mandatory to implement this method

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/DilatedConvED.py

@@ -228,7 +228,8 @@ class DilatedConvED(Base):
                                                 nn.BatchNorm1d(hidden_RNN) if use_bn else nn.Dropout(dropout_rate) ,
                                                 Permute() if use_bn else nn.Identity() ,
                                                 nn.Linear(hidden_RNN ,self.mul))
-        
+    def can_be_compiled(self):
+        return True  
 
 
     def forward(self, batch):

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/Duet.py

@@ -136,7 +136,8 @@ class Duet(Base):
                                          activation(),
                                          nn.Linear(dim*2,self.out_channels*self.mul  ))
 
-
+    def can_be_compiled(self):
+        return False  
     def forward(self, batch:dict)-> float:
         # x: [Batch, Input length, Channel]
         x_enc = batch['x_num_past'].to(self.device)

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/ITransformer.py

@@ -8,6 +8,8 @@ import numpy as np
 from .itransformer.Transformer_EncDec import Encoder, EncoderLayer
 from .itransformer.SelfAttention_Family import FullAttention, AttentionLayer
 from .itransformer.Embed import DataEmbedding_inverted
+from ..data_structure.utils import beauty_string
+from .utils import  get_scope,get_activation,Embedding_cat_variables
 
 try:
     import lightning.pytorch as pl
@@ -17,12 +19,6 @@ except:
     import pytorch_lightning as pl
     OLD_PL = True
     from .base import Base
-from .utils import QuantileLossMO,Permute, get_activation
-
-from typing import List, Union
-from ..data_structure.utils import beauty_string
-from .utils import  get_scope
-from .utils import Embedding_cat_variables
 
 
 
@@ -34,8 +30,6 @@ class ITransformer(Base):
     description = get_scope(handle_multivariate,handle_future_covariates,handle_categorical_variables,handle_quantile_loss)
     
     def __init__(self, 
-                
-
                  # specific params
                  hidden_size:int,
                  d_model: int,
@@ -107,6 +101,9 @@ class ITransformer(Base):
         )
         self.projector = nn.Linear(d_model, self.future_steps*self.mul, bias=True)
 
+    def can_be_compiled(self):
+        return True  
+
     def forecast(self, x_enc, x_mark_enc, x_dec, x_mark_dec):
         if self.use_norm:
             # Normalization from Non-stationary Transformer

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/Informer.py

@@ -124,7 +124,8 @@ class Informer(Base):
         
                 
         
-  
+    def can_be_compiled(self):
+        return True  
         
     def forward(self,batch): 
         #x_enc, x_mark_enc, x_dec, x_mark_dec,enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/LinearTS.py

@@ -143,7 +143,8 @@ class LinearTS(Base):
                                                     activation(),
                                                     nn.BatchNorm1d(hidden_size//8) if use_bn else nn.Dropout(dropout_rate) ,    
                                                     nn.Linear(hidden_size//8,self.future_steps*self.mul)))
-                               
+    def can_be_compiled(self):
+        return True                            
     def forward(self, batch):
       
         x =  batch['x_num_past'].to(self.device)

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/PatchTST.py

@@ -133,6 +133,9 @@ class PatchTST(Base):
     
         #self.final_linear = nn.Sequential(nn.Linear(past_channels,past_channels//2),activation(),nn.Dropout(dropout_rate), nn.Linear(past_channels//2,out_channels)  )
     
+    def can_be_compiled(self):
+        return True  
+    
     def forward(self, batch):           # x: [Batch, Input length, Channel]
         
 

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/RNN.py

@@ -148,7 +148,8 @@ class RNN(Base):
                                             activation(),
                                             MyBN(hidden_RNN//8) if use_bn else nn.Dropout(dropout_rate) ,
                                             nn.Linear(hidden_RNN//8,1)))
-
+    def can_be_compiled(self):
+        return True  
   
 
     def forward(self, batch):

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/Samformer.py

@@ -85,7 +85,9 @@ class Samformer(Base):
                                          activation(),
                                          nn.Linear(dim*2,self.out_channels*self.mul  ))
 
-
+    def can_be_compiled(self):
+        return True  
+  
     def forward(self, batch:dict)-> float:
 
         x = batch['x_num_past'].to(self.device)

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/Simple.py

@@ -67,7 +67,9 @@ class Simple(Base):
         self.linear = (nn.Sequential(nn.Linear(emb_past_out_channel*self.past_steps+emb_fut_out_channel*self.future_steps+self.past_steps*self.past_channels+self.future_channels*self.future_steps,hidden_size),
                                                     activation(),nn.Dropout(dropout_rate),
                                                     nn.Linear(hidden_size,self.out_channels*self.future_steps*self.mul)))
-                               
+    def can_be_compiled(self):
+        return True  
+                             
     def forward(self, batch):
       
         x =  batch['x_num_past'].to(self.device)

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/TFT.py

@@ -111,6 +111,10 @@ class TFT(Base):
 
         self.outLinear = nn.Linear(d_model, self.out_channels*self.mul)
 
+    def can_be_compiled(self):
+        return False  
+  
+
     def forward(self, batch:dict) -> torch.Tensor:
         """Temporal Fusion Transformer
 

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/TIDE.py

@@ -106,7 +106,10 @@ class TIDE(Base):
 
         # linear for Y lookback
         self.linear_target = nn.Linear(self.past_steps*self.out_channels, self.future_steps*self.out_channels*self.mul)
-
+    
+    def can_be_compiled(self):
+        return False  
+  
 
     def forward(self, batch:dict)-> float:
         """training process of the diffusion network

dsipts-1.1.12/src/dsipts/models/TTM.py

@@ -0,0 +1,158 @@
+import torch
+import numpy as np
+from torch import  nn
+
+try:
+    import lightning.pytorch as pl
+    from .base_v2 import Base
+    OLD_PL = False
+except:
+    import pytorch_lightning as pl
+    OLD_PL = True
+    from .base import Base
+
+
+from .ttm.utils import get_model, get_frequency_token, count_parameters, DEFAULT_FREQUENCY_MAPPING
+from ..data_structure.utils import beauty_string
+from .utils import  get_scope
+
+class TTM(Base):
+    handle_multivariate = True
+    handle_future_covariates = True
+    handle_categorical_variables = True
+    handle_quantile_loss = True
+    description = get_scope(handle_multivariate,handle_future_covariates,handle_categorical_variables,handle_quantile_loss)
+    
+    def __init__(self, 
+                model_path:str,
+                prefer_l1_loss:bool,  # exog: set true to use l1 loss
+                prefer_longer_context:bool,
+                prediction_channel_indices,
+                exogenous_channel_indices_cont,
+                exogenous_channel_indices_cat,
+                decoder_mode,
+                freq,
+                freq_prefix_tuning,
+                fcm_context_length,
+                fcm_use_mixer,
+                fcm_mix_layers,
+                fcm_prepend_past,
+                enable_forecast_channel_mixing,
+                **kwargs)->None:
+   
+        super().__init__(**kwargs)
+        self.save_hyperparameters(logger=False)
+
+        
+
+        self.index_fut = list(exogenous_channel_indices_cont)
+
+        if len(exogenous_channel_indices_cat)>0:
+            self.index_fut_cat = (self.past_channels+len(self.embs_past))+list(exogenous_channel_indices_cat)
+        else:
+            self.index_fut_cat = []
+        self.freq = freq
+
+        base_freq_token = get_frequency_token(self.freq)  # e.g., shape [n_token] or scalar
+        # ensure it's a tensor of integer type
+        if not torch.is_tensor(base_freq_token):
+            base_freq_token = torch.tensor(base_freq_token)
+        base_freq_token = base_freq_token.long()
+        self.register_buffer("token", base_freq_token, persistent=True)
+                
+        
+        self.model = get_model(
+            model_path=model_path,
+            context_length=self.past_steps,
+            prediction_length=self.future_steps,
+            prefer_l1_loss=prefer_l1_loss,
+            prefer_longer_context=prefer_longer_context,
+            num_input_channels=self.past_channels+len(self.embs_past), #giusto
+            decoder_mode=decoder_mode,
+            prediction_channel_indices=list(prediction_channel_indices),
+            exogenous_channel_indices=self.index_fut + self.index_fut_cat,
+            fcm_context_length=fcm_context_length,
+            fcm_use_mixer=fcm_use_mixer,
+            fcm_mix_layers=fcm_mix_layers,
+            freq=freq,
+            freq_prefix_tuning=freq_prefix_tuning,
+            fcm_prepend_past=fcm_prepend_past,
+            enable_forecast_channel_mixing=enable_forecast_channel_mixing,
+            
+        )
+        hidden_size =  self.model.config.hidden_size
+        self.model.prediction_head = torch.nn.Linear(hidden_size, self.out_channels*self.mul)
+        self._freeze_backbone()
+
+    def _freeze_backbone(self):
+        """
+        Freeze the backbone of the model.
+        This is useful when you want to fine-tune only the head of the model.
+        """
+        beauty_string(f"Number of params before freezing backbone:{count_parameters(self.model)}",'info',self.verbose)
+        
+        # Freeze the backbone of the model
+        for param in self.model.backbone.parameters():
+            param.requires_grad = False
+        # Count params
+        beauty_string(f"Number of params after freezing the backbone: {count_parameters(self.model)}",'info',self.verbose)
+            
+    
+    def _scaler_past(self, input):
+        for i, e in enumerate(self.embs_past):
+            input[:,:,i] = input[:, :, i] / (e-1)
+        return input
+    def _scaler_fut(self, input):
+        for i, e in enumerate(self.embs_fut):
+            input[:,:,i] = input[:, :, i] / (e-1)
+        return input
+
+    def can_be_compiled(self):
+        return True  
+  
+    def forward(self, batch):
+        x_enc = batch['x_num_past'].to(self.device)
+        original_indexes = batch['idx_target'][0].tolist()
+
+
+        if 'x_cat_past' in batch.keys():
+            x_mark_enc = batch['x_cat_past'].to(torch.float32).to(self.device)
+            x_mark_enc = self._scaler_past(x_mark_enc)
+            past_values = torch.cat((x_enc,x_mark_enc), axis=-1).type(torch.float32)
+        else:
+            past_values = x_enc
+        
+        future_values = torch.zeros_like(past_values).to(self.device)
+        future_values = future_values[:,:self.future_steps,:]
+
+        if 'x_num_future' in batch.keys(): 
+            future_values[:,:,self.index_fut] = batch['x_num_future'].to(self.device)
+        if 'x_cat_future' in batch.keys():
+            x_mark_dec = batch['x_cat_future'].to(torch.float32).to(self.device)
+            x_mark_dec = self._scaler_fut(x_mark_dec)
+            future_values[:,:,self.index_cat_fut] = x_mark_dec
+        
+
+        #investigating!! problem with dynamo!
+        #freq_token = get_frequency_token(self.freq).repeat(past_values.shape[0])
+
+        batch_size = past_values.shape[0]
+        freq_token = self.token.repeat(batch_size).long().to(self.device)
+
+
+        res = self.model(
+            past_values= past_values,
+            future_values= future_values,# future_values if future_values.shape[0]>0 else None,
+            past_observed_mask = None,
+            future_observed_mask = None,
+            output_hidden_states =  False,
+            return_dict = False,
+            freq_token= freq_token,#[0:past_values.shape[0]], ##investigating
+            static_categorical_values = None
+        )
+
+
+        BS = res.shape[0]
+        return res.reshape(BS,self.future_steps,-1,self.mul)
+        
+    

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/TimeXER.py

@@ -125,7 +125,9 @@ class TimeXER(Base):
 
         
 
-
+    def can_be_compiled(self):
+        return True  
+  
 
 
     def forward(self, batch:dict)-> float:

{dsipts-1.1.10 → dsipts-1.1.12}/src/dsipts/models/base.py

@@ -111,8 +111,11 @@ class Base(pl.LightningModule):
         self.train_loss_epoch = -100.0
         self.verbose = verbose
         self.name = self.__class__.__name__
-        self.train_epoch_metrics = []
-        self.validation_epoch_metrics = []
+        self.register_buffer("train_epoch_metrics", torch.tensor(0.0))
+        self.register_buffer("validation_epoch_metrics", torch.tensor(0.0))
+        self.register_buffer("train_epoch_count", torch.tensor(0))
+        self.register_buffer("validation_epoch_count", torch.tensor(0))
+        
         
         self.use_quantiles = True if len(quantiles)>0 else False
         self.quantiles =  quantiles
@@ -295,7 +298,8 @@ class Base(pl.LightningModule):
                 y_hat = self(batch)
                 loss = self.compute_loss(batch,y_hat)
             
-        self.train_epoch_metrics.append(loss.item())
+        self.train_epoch_metrics+=loss.detach()
+        self.train_epoch_count +=1
         return loss
 
     
@@ -311,27 +315,20 @@ class Base(pl.LightningModule):
             y_hat = self(batch)
             score = 0
         if batch_idx==0:
-            if self.use_quantiles:
-                idx = 1
-            else:
-                idx = 0
-            #track the predictions! We can do better than this but maybe it is better to firstly update pytorch-lightening 
-
+           
             if self.count_epoch%int(max(self.trainer.max_epochs/100,1))==1:
-
-                for i in range(batch['y'].shape[2]):
-                    real =  batch['y'][0,:,i].cpu().detach().numpy()
-                    pred =  y_hat[0,:,i,idx].cpu().detach().numpy()
-                    fig, ax = plt.subplots(figsize=(7,5))  
-                    ax.plot(real,'o-',label='real')
-                    ax.plot(pred,'o-',label='pred')
-                    ax.legend()
-                    ax.set_title(f'Channel {i} first element first batch validation {int(100*self.count_epoch/self.trainer.max_epochs)}%')
-                    self.logger.experiment.track(Image(fig), name='cm_training_end')
-                    #self.log(f"example_{i}", np.stack([real, pred]).T,sync_dist=True)
-
-        return self.compute_loss(batch,y_hat)+score
-
+                self._val_outputs.append({
+                    "y": batch['y'].detach().cpu(),
+                    "y_hat": y_hat.detach().cpu()
+                })                
+        self.validation_epoch_metrics = (self.compute_loss(batch,y_hat)+score).detach()
+        self.validation_epoch_count+=1
+
+        return None #self.compute_loss(batch,y_hat)+score
+    
+    def on_validation_start(self):
+        # reset buffer each epoch
+        self._val_outputs = []
 
     def validation_epoch_end(self, outs):
         """
@@ -339,14 +336,30 @@ class Base(pl.LightningModule):
         
         :meta private:
         """
-        if len(outs)==0:
-            loss = 10000
-            beauty_string(f'THIS IS A BUG, It should be polulated','info',self.verbose)
-        else:
-            loss = torch.stack(outs).mean()
-        
-        self.log("val_loss", loss.item(),sync_dist=True)
-        beauty_string(f'Epoch: {self.count_epoch} train error: {self.train_loss_epoch:.4f} validation loss: {loss.item():.4f}','info',self.verbose)
+        if len(self._val_outputs)>0:
+            ys = torch.cat([o["y"] for o in self._val_outputs])
+            y_hats = torch.cat([o["y_hat"] for o in self._val_outputs])
+            if self.use_quantiles:
+                idx = 1
+            else:
+                idx = 0
+            for i in range(ys.shape[2]):
+                real =  ys[0,:,i].cpu().detach().numpy()
+                pred =  y_hats[0,:,i,idx].cpu().detach().numpy()
+                fig, ax = plt.subplots(figsize=(7,5))  
+                ax.plot(real,'o-',label='real')
+                ax.plot(pred,'o-',label='pred')
+                ax.legend()
+                ax.set_title(f'Channel {i} first element first batch validation {int(100*self.count_epoch/self.trainer.max_epochs)}%')
+                self.logger.experiment.track(Image(fig), name='cm_training_end')
+                #self.log(f"example_{i}", np.stack([real, pred]).T,sync_dist=True)
+                plt.close(fig) 
+        avg = self.validation_epoch_metrics/self.validation_epoch_count
+
+        self.validation_epoch_metrics.zero_()
+        self.validation_epoch_count.zero_()
+        self.log("val_loss", avg,sync_dist=True)
+        beauty_string(f'Epoch: {self.count_epoch} train error: {self.train_loss_epoch:.4f} validation loss: {avg:.4f}','info',self.verbose)
 
     def training_epoch_end(self, outs):
         """
@@ -355,12 +368,11 @@ class Base(pl.LightningModule):
         :meta private:
         """
 
-        loss = sum(outs['loss'] for outs in outs) / len(outs)
-        self.log("train_loss", loss.item(),sync_dist=True)
+        loss = self.train_epoch_metrics/self.global_step
+        self.log("train_loss", loss,sync_dist=True)
         self.count_epoch+=1    
 
-        self.train_loss_epoch = loss.item()
-
+        self.train_loss_epoch = loss
     def compute_loss(self,batch,y_hat):
         """
         custom loss calculation