No need to use the softmax layer since CrossEntropyLoss already uses …

…it internally (issue #1).
vinhkhuc · Feb 26, 2017 · f32192e · f32192e
1 parent 4ee3ccb
commit f32192e
Show file tree

Hide file tree

Showing 4 changed files with 3 additions and 7 deletions.
diff --git a/2_logistic_regression.py b/2_logistic_regression.py
@@ -8,9 +8,9 @@
 
 
 def build_model(input_dim, output_dim):
+    # We don't need the softmax layer here since CrossEntropyLoss already uses it internally.
     model = torch.nn.Sequential()
     model.add_module("linear", torch.nn.Linear(input_dim, output_dim, bias=False))
-    model.add_module("softmax", torch.nn.Softmax())
     return model
 
 

diff --git a/3_neural_net.py b/3_neural_net.py
@@ -12,7 +12,6 @@ def build_model(input_dim, output_dim):
     model.add_module("linear_1", torch.nn.Linear(input_dim, 512, bias=False))
     model.add_module("sigmoid_1", torch.nn.Sigmoid())
     model.add_module("linear_2", torch.nn.Linear(512, output_dim, bias=False))
-    model.add_module("softmax", torch.nn.Softmax())
     return model
 
 

diff --git a/4_modern_neural_net.py b/4_modern_neural_net.py
@@ -16,7 +16,6 @@ def build_model(input_dim, output_dim):
     model.add_module("relu_2", torch.nn.ReLU())
     model.add_module("dropout_2", torch.nn.Dropout(0.2))
     model.add_module("linear_3", torch.nn.Linear(512, output_dim, bias=False))
-    model.add_module("softmax", torch.nn.Softmax())
     return model
 
 

diff --git a/5_convolutional_net.py b/5_convolutional_net.py
@@ -7,7 +7,7 @@
 from data_util import load_mnist
 
 
-# Separately create two sequential here since PyTorch doesn't have nn.View()
+# We need to create two sequential models here since PyTorch doesn't have nn.View()
 class ConvNet(torch.nn.Module):
     def __init__(self, output_dim):
         super(ConvNet, self).__init__()
@@ -26,8 +26,6 @@ def __init__(self, output_dim):
         self.fc.add_module("relu_3", torch.nn.ReLU())
         self.fc.add_module("dropout_3", torch.nn.Dropout())
         self.fc.add_module("fc2", torch.nn.Linear(50, output_dim))
-        self.fc.add_module("relu_4", torch.nn.ReLU())
-        self.fc.add_module("softmax", torch.nn.Softmax())
 
     def forward(self, x):
         x = self.conv.forward(x)
@@ -75,7 +73,7 @@ def main():
     n_classes = 10
     model = ConvNet(output_dim=n_classes)
     loss = torch.nn.CrossEntropyLoss(size_average=True)
-    optimizer = optim.SGD(model.parameters())
+    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
     batch_size = 100
 
     for i in range(100):