project Refactor, CNN Classifier Basics

lib/models/homotopy_classification/cnn_based.py

@@ -24,41 +24,44 @@ class ConvHomDetector(LightningBaseModule):
     def training_step(self, batch_xy, batch_nb, *args, **kwargs):
         batch_x, batch_y = batch_xy
         pred_y = self(batch_x)
-        loss = F.binary_cross_entropy(pred_y, batch_y.float())
+        loss = self.criterion(pred_y, batch_y.unsqueeze(-1).float())
         return {'loss': loss, 'log': dict(loss=loss)}
 
-    def test_step(self, batch_xy, **kwargs):
+    def test_step(self, batch_xy, batch_nb, **kwargs):
         batch_x, batch_y = batch_xy
         pred_y = self(batch_x)
-        return dict(prediction=pred_y, label=batch_y)
+        return dict(prediction=pred_y, label=batch_y, batch_nb=batch_nb)
 
-    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])
+    def test_epoch_end(self, outputs):
+        evaluation = ROCEvaluation(plot_roc=True)
+        predictions = torch.cat([x['prediction'] for x in outputs])
+        labels = torch.cat([x['label'] for x in outputs]).unsqueeze(1)
 
-        scores = evaluation(predictions.numpy(), labels.numpy(), )
-        self.logger.log_metrics({key:value for key, value in zip(['roc_auc', 'tpr', 'fpr'], scores)})
+        # Sci-py call ROC eval call is eval(true_label, prediction)
+        roc_auc, tpr, fpr = evaluation(labels.cpu().numpy(), predictions.cpu().numpy(), )
+        score_dict = dict(roc_auc=roc_auc)
+        # self.logger.log_metrics(score_dict)
         self.logger.log_image(f'{self.name}', plt.gcf())
-        pass
 
-    def __init__(self, *params):
-        super(ConvHomDetector, self).__init__(*params)
+        return dict(log=score_dict)
+
+    def __init__(self, hparams):
+        super(ConvHomDetector, self).__init__(hparams)
 
         # Dataset
-        self.dataset = TrajData(self.hparams.data_param.root)
+        self.dataset = TrajData(self.hparams.data_param.map_root, mode='all_in_map')
 
         # Additional Attributes
         self.map_shape = self.dataset.map_shapes_max
 
-        # Model Paramters
+        # Model Parameters
         self.in_shape = self.dataset.map_shapes_max
         assert len(self.in_shape) == 3, f'Image or map shape has to have 3 dims, but had: {len(self.in_shape)}'
+        self.criterion = nn.BCEWithLogitsLoss()
 
         # NN Nodes
         # ============================
         # Convolutional Map Processing
-        #
         self.map_conv_0 = ConvModule(self.in_shape, conv_kernel=3, conv_stride=1,
                                      conv_padding=0, conv_filters=self.hparams.model_param.filters[0])
         self.map_res_1 = ResidualModule(self.map_conv_0.shape, ConvModule, 3,
@@ -86,7 +89,6 @@ class ConvHomDetector(LightningBaseModule):
         self.linear = nn.Linear(reduce(mul, self.flatten.shape), self.hparams.model_param.classes * 10)
         # Comments on Multi Class labels
         self.classifier = nn.Linear(self.hparams.model_param.classes * 10, 1)  # self.hparams.model_param.classes)
-        self.out_activation = nn.Sigmoid()  # nn.Softmax
 
     def forward(self, x):
         tensor = self.map_conv_0(x)
@@ -98,25 +100,4 @@ class ConvHomDetector(LightningBaseModule):
         tensor = self.flatten(tensor)
         tensor = self.linear(tensor)
         tensor = self.classifier(tensor)
-        tensor = self.out_activation(tensor)
         return tensor
-
-    # 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)