eval written

lib/models/homotopy_classification/cnn_based.py

@@ -5,8 +5,10 @@ import torch
 from torch import nn
 import torch.nn.functional as F
 from torch.optim import Adam
+from torch.utils.data import DataLoader
 
 from datasets.trajectory_dataset import TrajData
+from lib.evaluation.classification import ROCEvaluation
 from lib.modules.utils import LightningBaseModule, Flatten
 from lib.modules.blocks import ConvModule, ResidualModule
 
@@ -24,6 +26,22 @@ class ConvHomDetector(LightningBaseModule):
         loss = F.binary_cross_entropy(pred_y, batch_y.float())
         return {'loss': loss, 'log': dict(loss=loss)}
 
+    def test_step(self, batch_xy, **kwargs):
+        batch_x, batch_y = batch_xy
+        pred_y = self(batch_x)
+        return dict(prediction=pred_y, label=batch_y)
+
+    def test_end(self, outputs):
+        evaluation = ROCEvaluation()
+        predictions = torch.stack([x['prediction'] for x in outputs])
+        labels = torch.stack([x['label'] for x in outputs])
+
+        scores = evaluation(predictions.numpy(), labels.numpy())
+        self.logger.log_metrics()
+
+
+        pass
+
     def __init__(self, *params):
         super(ConvHomDetector, self).__init__(*params)
 
@@ -70,6 +88,26 @@ class ConvHomDetector(LightningBaseModule):
         self.classifier = nn.Linear(self.hparams.model_param.classes * 10, 1)  # self.hparams.model_param.classes)
         self.out_activation = nn.Sigmoid()  # nn.Softmax
 
+    # Dataloaders
+    # ================================================================================
+    # Train Dataloader
+    def train_dataloader(self):
+        return DataLoader(dataset=self.dataset.train_dataset, shuffle=True,
+                          batch_size=self.hparams.data_param.batchsize,
+                          num_workers=self.hparams.data_param.worker)
+
+    # Test Dataloader
+    def test_dataloader(self):
+        return DataLoader(dataset=self.dataset.test_dataset, shuffle=True,
+                          batch_size=self.hparams.data_param.batchsize,
+                          num_workers=self.hparams.data_param.worker)
+
+    # Validation Dataloader
+    def val_dataloader(self):
+        return DataLoader(dataset=self.dataset.val_dataset, shuffle=True,
+                          batch_size=self.hparams.data_param.batchsize,
+                          num_workers=self.hparams.data_param.worker)
+
     def forward(self, x):
         tensor = self.map_conv_0(x)
         tensor = self.map_res_1(tensor)