Variational Generator

lib/models/homotopy_classification/cnn_based.py

@@ -32,24 +32,35 @@ class ConvHomDetector(LightningBaseModule):
         pred_y = self(batch_x)
         return dict(prediction=pred_y, label=batch_y, batch_nb=batch_nb)
 
+    def validation_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, batch_nb=batch_nb)
+
     def test_epoch_end(self, outputs):
-        evaluation = ROCEvaluation(plot_roc=True)
+        return self._val_test_end(outputs)
+
+    def validation_epoch_end(self, outputs: list):
+        return self._val_test_end(outputs)
+
+    def _val_test_end(self, outputs, test=True):
+        evaluation = ROCEvaluation(plot_roc=True if test else False)
         predictions = torch.cat([x['prediction'] for x in outputs])
         labels = torch.cat([x['label'] for x in outputs]).unsqueeze(1)
 
         # 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)
+        roc_auc, tpr, fpr = evaluation(labels.cpu().numpy(), predictions.cpu().numpy())
         # self.logger.log_metrics(score_dict)
-        self.logger.log_image(f'{self.name}', plt.gcf())
+        if test:
+            self.logger.log_image(f'{self.name}', plt.gcf())
 
-        return dict(log=score_dict)
+        return dict(score=roc_auc, log=dict(roc_auc=roc_auc))
 
     def __init__(self, hparams):
         super(ConvHomDetector, self).__init__(hparams)
 
         # Dataset
-        self.dataset = TrajData(self.hparams.data_param.map_root, mode='all_in_map')
+        self.dataset = TrajData(self.hparams.data_param.map_root, mode='all_in_map', )
 
         # Additional Attributes
         self.map_shape = self.dataset.map_shapes_max
@@ -59,6 +70,7 @@ class ConvHomDetector(LightningBaseModule):
         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.BCELoss()
         self.sigmoid = nn.Sigmoid()
+        self.relu = nn.ReLU()
 
         # NN Nodes
         # ============================
@@ -100,6 +112,7 @@ class ConvHomDetector(LightningBaseModule):
         tensor = self.map_conv_3(tensor)
         tensor = self.flatten(tensor)
         tensor = self.linear(tensor)
+        tensor = self.relu(tensor)
         tensor = self.classifier(tensor)
         tensor = self.sigmoid(tensor)
         return tensor