Dataset rdy

util/module_mixins.py

@@ -0,0 +1,95 @@
+from abc import ABC
+
+import torch
+
+from ml_lib.modules.util import LightningBaseModule
+from util.loss_mixin import LossMixin
+from util.optimizer_mixin import OptimizerMixin
+
+
+class TrainMixin:
+
+    def training_step(self, batch_xy, batch_nb, *args, **kwargs):
+        assert isinstance(self, LightningBaseModule)
+        batch_x, batch_y = batch_xy
+        y = self(batch_x).main_out
+        loss = self.ce_loss(y.squeeze(), batch_y.long())
+        return dict(loss=loss)
+
+    def training_epoch_end(self, outputs):
+        assert isinstance(self, LightningBaseModule)
+        keys = list(outputs[0].keys())
+
+        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])
+
+
+class ValMixin:
+
+    def validation_step(self, batch_xy, batch_idx, *args, **kwargs):
+        assert isinstance(self, LightningBaseModule)
+        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())
+
+        return dict(val_loss=val_loss,
+                    batch_idx=batch_idx, y=y, batch_y=batch_y)
+
+    def validation_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}
+                            )
+
+        additional_scores = self.additional_scores(outputs)
+        summary_dict.update(**additional_scores)
+
+        for key in summary_dict.keys():
+            self.log(key, summary_dict[key])
+
+
+class TestMixin:
+
+    def test_step(self, batch_xy, batch_idx, *_, **__):
+        assert isinstance(self, LightningBaseModule)
+        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)
+
+    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}
+                            )
+
+        additional_scores = self.additional_scores(outputs)
+        summary_dict.update(**additional_scores)
+
+        for key in summary_dict.keys():
+            self.log(key, summary_dict[key])
+
+
+class CombinedModelMixins(LossMixin,
+                          TrainMixin,
+                          ValMixin,
+                          TestMixin,
+                          OptimizerMixin,
+                          LightningBaseModule,
+                          ABC):
+    pass