WIP: Lightning

Working on a training mechanism using pytorch lightning

Still a WIP

bob/learn/pytorch/architectures/lenet.py

@@ -5,7 +5,7 @@ from torch import nn
 
 
 class Lenet5(Module):
-    def __init__(self, num_features=30):
+    def __init__(self, num_features=84):
         super(Lenet5, self).__init__()
         self.conv1 = nn.Conv2d(1, 6, 5)
         self.relu1 = nn.ReLU()
@@ -15,12 +15,12 @@ class Lenet5(Module):
         self.pool2 = nn.MaxPool2d(2)
         self.fc1 = nn.Linear(256, 120)
         self.relu3 = nn.ReLU()
-        self.fc2 = nn.Linear(120, 84)
+        self.fc2 = nn.Linear(120, num_features)
         self.relu4 = nn.ReLU()
         # Tiago: this was a slight adaptation I made
         # self.fc3 = nn.Linear(84, 10)
-        self.fc3 = nn.Linear(84, num_features)
-        self.relu5 = nn.ReLU()
+        # self.fc3 = nn.Linear(84, num_features)
+        # self.relu5 = nn.ReLU()
 
     def forward(self, x):
 
@@ -35,6 +35,6 @@ class Lenet5(Module):
         y = self.relu3(y)
         y = self.fc2(y)
         y = self.relu4(y)
-        y = self.fc3(y)
-        y = self.relu5(y)
+        # y = self.fc3(y)
+        # y = self.relu5(y)
         return y