Model Init

models/binary_classifier.py

@@ -0,0 +1,47 @@
+import torch
+from torch import nn
+from torch.optim import Adam
+
+from ml_lib.modules.blocks import ConvModule
+from ml_lib.modules.utils import LightningBaseModule
+
+
+class BinaryClassifier(LightningBaseModule):
+    @classmethod
+    def name(cls):
+        return cls.__name__
+
+    def configure_optimizers(self):
+        return Adam(lr=self.hparams.train.lr)
+
+    def training_step(self, batch_xy, batch_nb, *args, **kwargs):
+        batch_x, batch_y = batch_xy
+        y = self(batch_y)
+        loss = self.criterion(y, batch_y)
+        return dict(loss=loss)
+
+    def validation_step(self, batch_xy, **kwargs):
+        batch_x, batch_y = batch_xy
+        y = self(batch_y)
+        val_loss = self.criterion(y, batch_y)
+        return dict(val_loss=val_loss)
+
+    def validation_epoch_end(self, outputs):
+        over_all_val_loss = torch.mean(torch.stack([output['val_loss'] for output in outputs]))
+
+    def __init__(self, hparams):
+        super(BinaryClassifier, self).__init__(hparams)
+        self.criterion = nn.BCELoss()
+
+        # Additional parameters
+        self.in_shape = ()
+
+        # Model Modules
+        self.conv_1 = ConvModule(self.in_shape, 32, 5, )
+        self.conv_2 = ConvModule(64)
+        self.conv_3 = ConvModule(128)
+
+    def forward(self, batch, **kwargs):
+        return batch
+
+