project Refactor, CNN Classifier Basics

lib/models/generators/full.py

@@ -1,14 +1,16 @@
-from lib.modules.blocks import LightningBaseModule
 from lib.modules.losses import BinaryHomotopicLoss
+from lib.modules.utils import LightningBaseModule
 from lib.objects.map import Map
 from lib.objects.trajectory import Trajectory
 
 import torch.nn as nn
 
-nn.MSELoss
 
 class LinearRouteGeneratorModel(LightningBaseModule):
 
+    def test_epoch_end(self, outputs):
+        pass
+
     name = 'LinearRouteGenerator'
 
     def configure_optimizers(self):
@@ -33,6 +35,12 @@ class LinearRouteGeneratorModel(LightningBaseModule):
         pred_y = self(map_x, traj_x, label_x)
 
         loss = self.loss(traj_x, pred_y)
+
+        def training_step(self, batch_xy, batch_nb, *args, **kwargs):
+            batch_x, batch_y = batch_xy
+            pred_y = self(batch_x)
+            loss = self.criterion(pred_y, batch_y.unsqueeze(-1).float())
+
         return dict(loss=loss, log=dict(loss=loss))
 
     def test_step(self, *args, **kwargs):
@@ -41,7 +49,7 @@ class LinearRouteGeneratorModel(LightningBaseModule):
     def __init__(self, *params):
         super(LinearRouteGeneratorModel, self).__init__(*params)
 
-        self.loss = BinaryHomotopicLoss(self.map_storage)
+        self.criterion = BinaryHomotopicLoss(self.map_storage)
 
     def forward(self, map_x, traj_x, label_x):
         pass