Initial Commit

src/trainers/autoregressive_trainer.py

@@ -0,0 +1,59 @@
+from clearml import OutputModel
+import torch
+from data.preprocessing import DataProcessor
+from utils.clearml import ClearMLHelper
+from utils.autoregressive import predict_auto_regressive
+import plotly.graph_objects as go
+import numpy as np
+import plotly.subplots as sp
+from plotly.subplots import make_subplots
+from trainers.trainer import Trainer
+
+class AutoRegressiveTrainer(Trainer):
+    def debug_plots(self, task, train: bool, samples, epoch):
+        X, y = samples
+        X = X.to(self.device)
+
+        num_samples = len(X)
+        rows = num_samples  # One row per sample since we only want one column
+        cols = 1
+        
+        fig = make_subplots(rows=rows, cols=cols, subplot_titles=[f'Sample {i+1}' for i in range(num_samples)])
+
+        for i, (current_day, next_day) in enumerate(zip(X, y)):
+            predictions = self.predict_auto_regressive(current_day)
+            sub_fig = self.get_plot(current_day, next_day, predictions, show_legend=(i == 0))
+            
+            row = i + 1
+            col = 1
+            
+            for trace in sub_fig.data:
+                fig.add_trace(trace, row=row, col=col)
+
+            loss = self.criterion(predictions.to(self.device), next_day.to(self.device)).item()
+
+            fig['layout']['annotations'][i].update(text=f"{loss.__class__.__name__}: {loss:.6f}")
+
+        # y axis same for all plots
+        fig.update_yaxes(range=[-1, 1], col=1)
+
+        
+        fig.update_layout(height=300 * rows)
+        task.get_logger().report_plotly(
+            title=f"{'Training' if train else 'Test'} Samples",
+            series="full_day",
+            iteration=epoch,
+            figure=fig
+        )
+
+
+    def predict_auto_regressive(self, initial_sequence: torch.Tensor, sequence_length: int = 96):
+        initial_sequence = initial_sequence.to(self.device)
+
+        return predict_auto_regressive(self.model, initial_sequence, sequence_length)
+
+    def random_day_prediction(self):
+        current_day_features, next_day_features = self.data_processor.get_random_test_day()
+
+        predictions = self.predict_auto_regressive(current_day_features)
+        return current_day_features, next_day_features, predictions