Transformer running

util/module_mixins.py

@@ -1,6 +1,7 @@
 from abc import ABC
 
 import torch
+import pandas as pd
 
 from ml_lib.modules.util import LightningBaseModule
 from util.loss_mixin import LossMixin
@@ -11,9 +12,15 @@ class TrainMixin:
 
     def training_step(self, batch_xy, batch_nb, *args, **kwargs):
         assert isinstance(self, LightningBaseModule)
-        batch_x, batch_y = batch_xy
+        batch_files, batch_x, batch_y = batch_xy
         y = self(batch_x).main_out
-        loss = self.ce_loss(y.squeeze(), batch_y.long())
+
+        if self.params.loss == 'focal_loss_rob':
+            labels_one_hot = torch.nn.functional.one_hot(batch_y, num_classes=5)
+            loss = self.__getattribute__(self.params.loss)(y, labels_one_hot)
+        else:
+            loss = self.__getattribute__(self.params.loss)(y, batch_y.long())
+
         return dict(loss=loss)
 
     def training_epoch_end(self, outputs):
@@ -23,21 +30,23 @@ class TrainMixin:
         summary_dict = {f'mean_{key}': torch.mean(torch.stack([output[key]
                                                                for output in outputs]))
                         for key in keys if 'loss' in key}
-        for key in summary_dict.keys():
-            self.log(key, summary_dict[key])
+        summary_dict.update(epoch=self.current_epoch)
+        self.log_dict(summary_dict)
 
 
 class ValMixin:
 
     def validation_step(self, batch_xy, batch_idx, *args, **kwargs):
         assert isinstance(self, LightningBaseModule)
-        batch_x, batch_y = batch_xy
+        batch_files, batch_x, batch_y = batch_xy
         model_out = self(batch_x)
         y = model_out.main_out
 
-        val_loss = self.ce_loss(y.squeeze(), batch_y.long())
+        val_loss = self.ce_loss(y, batch_y.long())
 
-        return dict(val_loss=val_loss,
+        self.metrics.update(y, batch_y)  # torch.argmax(y, -1), batch_y)
+
+        return dict(val_loss=val_loss, batch_files=batch_files,
                     batch_idx=batch_idx, y=y, batch_y=batch_y)
 
     def validation_epoch_end(self, outputs, *_, **__):
@@ -49,40 +58,40 @@ class ValMixin:
                                                                     for output in outputs]))
                              for key in keys if 'loss' in key}
                             )
-
+        # Sklearn Scores
         additional_scores = self.additional_scores(outputs)
         summary_dict.update(**additional_scores)
 
-        for key in summary_dict.keys():
-            self.log(key, summary_dict[key])
+        pl_metrics, pl_images = self.metrics.compute_and_prepare()
+        self.metrics.reset()
+        summary_dict.update(**pl_metrics)
+        summary_dict.update(epoch=self.current_epoch)
+
+        self.log_dict(summary_dict, on_epoch=True)
+
+        for name, image in pl_images.items():
+            self.logger.log_image(name, image, step=self.global_step)
+        pass
 
 
 class TestMixin:
 
     def test_step(self, batch_xy, batch_idx, *_, **__):
         assert isinstance(self, LightningBaseModule)
-        batch_x, batch_y = batch_xy
+        batch_files, batch_x, batch_y = batch_xy
         model_out = self(batch_x)
         y = model_out.main_out
-        test_loss = self.ce_loss(y.squeeze(), batch_y.long())
-        return dict(test_loss=test_loss,
-                    batch_idx=batch_idx, y=y, batch_y=batch_y)
+        return dict(batch_files=batch_files, batch_idx=batch_idx, y=y)
 
     def test_epoch_end(self, outputs, *_, **__):
         assert isinstance(self, LightningBaseModule)
-        summary_dict = dict()
 
-        keys = list(outputs[0].keys())
-        summary_dict.update({f'mean_{key}': torch.mean(torch.stack([output[key]
-                                                                    for output in outputs]))
-                             for key in keys if 'loss' in key}
-                            )
+        y_arg_max = torch.argmax(outputs[0]['y'])
 
-        additional_scores = self.additional_scores(outputs)
-        summary_dict.update(**additional_scores)
+        pd.DataFrame(data=dict(filenames=outputs[0]['batch_files'], predtiction=y_arg_max))
 
-        for key in summary_dict.keys():
-            self.log(key, summary_dict[key])
+        # No logging, just inference.
+        # self.log_dict(summary_dict, on_epoch=True)
 
 
 class CombinedModelMixins(LossMixin,