Removing softmax, closes #37

tests/test_models.py

@@ -17,9 +17,6 @@ def test_christian_model(image_shape, num_classes):
     y = model(x)
 
     assert y.shape == (n, num_classes), f"Shape: {y.shape}"
-    assert y.sum(dim=1).allclose(torch.ones(n), atol=1e-5), (
-        f"Softmax output should sum to 1, but got: {y.sum()}"
-    )
 
 
 @pytest.mark.parametrize(
@@ -35,3 +32,4 @@ def test_jan_model(image_shape, num_classes):
     y = model(x)
 
     assert y.shape == (n, num_classes), f"Shape: {y.shape}"
+

utils/models/christian_model.py

@@ -58,15 +58,13 @@ def __init__(self, image_shape, num_classes):
         self.cnn2 = CNNBlock(50, 100)
 
         self.fc1 = nn.Linear(100 * 4 * 4, num_classes)
-        self.softmax = nn.Softmax(dim=1)
 
     def forward(self, x):
         x = self.cnn1(x)
         x = self.cnn2(x)
 
         x = x.view(x.size(0), -1)
         x = self.fc1(x)
-        x = self.softmax(x)
 
         return x
 

utils/models/johan_model.py

@@ -43,14 +43,12 @@ def __init__(self, image_shape, num_classes):
         self.fc2 = nn.Linear(77, 77)
         self.fc3 = nn.Linear(77, 77)
         self.fc4 = nn.Linear(77, num_classes)
-        self.softmax = nn.Softmax(dim=1)
         self.relu = nn.ReLU()
 
     def forward(self, x):
         for layer in [self.fc1, self.fc2, self.fc3, self.fc4]:
             x = layer(x)
             x = self.relu(x)
-        x = self.softmax(x)
         return x
 
 

utils/models/solveig_model.py

@@ -58,7 +58,6 @@ def forward(self, x):
         x = torch.flatten(x, 1)
 
         x = self.fc1(x)
-        x = nn.Softmax(x)
 
         return x