change net

xiyuan27 · Aug 10, 2021 · 84fcee1 · 84fcee1
1 parent 308e11f
commit 84fcee1
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 36 deletions.
diff --git a/torchTrain.py b/torchTrain.py
@@ -6,6 +6,7 @@
 from torch.autograd import Variable
 import torch.optim as optim
 import time
+from torchvision import models
 
 
 DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -16,8 +17,8 @@ def loadtraindata(path):
     # print('path:', path)
     # print('filename:', filename)
     trainset = torchvision.datasets.ImageFolder(path,
-                transform=transforms.Compose([transforms.Resize((32, 32)),  # resize image (h,w)
-                transforms.CenterCrop(32), transforms.ToTensor()]))
+                transform=transforms.Compose([transforms.Resize((224, 224)),  # resize image (h,w)
+                transforms.CenterCrop(224), transforms.ToTensor()]))
     print('classes: ', trainset.classes)  # all classes in dataset
     print('number of classes: ', len(trainset.classes))
 
@@ -42,7 +43,7 @@ def forward(self, x):  # feed forward
 
         x = self.pool(F.relu(self.conv1(x)))  # F is torch.nn.functional
         x = self.pool(F.relu(self.conv2(x)))
-        x = x.view(-1, 16 * 5 * 5)  # .view( ) is a method tensor, which automatically change tensor size but elements number not change
+        x = x.view(x.shape[0], -1)  # .view( ) is a method tensor, which automatically change tensor size but elements number not change
 
         x = F.relu(self.fc1(x))
         x = F.relu(self.fc2(x))
@@ -54,10 +55,23 @@ def forward(self, x):  # feed forward
 def trainandsave(path):
     trainloader, filename, classes = loadtraindata(path)
 
-    net = Net(len(classes))
+    # net = Net(len(classes))
+    # print(net)
+
+    net = models.alexnet(pretrained=True)
+    net.classifier = nn.Sequential(
+        nn.Dropout(),
+        nn.Linear(256 * 6 * 6, 4096),
+        nn.ReLU(inplace=True),
+        nn.Dropout(),
+        nn.Linear(4096, 4096),
+        nn.ReLU(inplace=True),
+        nn.Linear(4096, len(classes)),
+    )
     print(net)
+
     net.to(DEVICE)
-    optimizer = optim.SGD(net.parameters(), lr=0.008, momentum=0.9)  # learning rate=0.002
+    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)  # learning rate=0.002
     criterion = nn.CrossEntropyLoss()  # loss function
     # training part
     for epoch in range(5):  # 5 epoch

diff --git a/torchTrain_test.py b/torchTrain_test.py
@@ -57,8 +57,7 @@ def forward(self, x):  # feed forward
         x = self.pool(F.relu(self.conv1(x)))  # F is torch.nn.functional
         x = self.pool(F.relu(self.conv2(x)))
         # print(x.shape)
-        x = x.view(x.shape[0],
-                   -1)  # .view( ) is a reshape operation, which automatically change tensor size but elements number not change
+        x = x.view(x.shape[0], -1)  # .view( ) is a reshape operation, which automatically change tensor size but elements number not change
         x = F.relu(self.fc1(x))
         x = F.relu(self.fc2(x))
         x = self.fc3(x)
@@ -67,37 +66,39 @@ def forward(self, x):  # feed forward
 
 def trainandsave(path):
     train_loader, validate_loader, test_loader, filename, classes = loadtraindata(path)
-
+    # network 1:
     # net = Net(len(classes))
-    # print(net)
-
-    net = models.alexnet(pretrained=True)
-    net.features = nn.Sequential(
-        nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
-        nn.ReLU(inplace=True),
-        nn.MaxPool2d(kernel_size=3, stride=2),
-        nn.Conv2d(64, 192, kernel_size=5, padding=2),
-        nn.ReLU(inplace=True),
-        nn.MaxPool2d(kernel_size=3, stride=2),
-        nn.Conv2d(192, 384, kernel_size=3, padding=1),
-        nn.ReLU(inplace=True),
-        nn.Conv2d(384, 256, kernel_size=3, padding=1),
-        nn.ReLU(inplace=True),
-        # nn.Conv2d(256, 256, kernel_size=3, padding=1),
-        # nn.ReLU(inplace=True),
-        nn.MaxPool2d(kernel_size=3, stride=2),
-    )
-    net.classifier = nn.Sequential(
-        nn.Dropout(),
-        nn.Linear(256 * 6 * 6, 4096),
-        nn.ReLU(inplace=True),
-        nn.Dropout(),
-        nn.Linear(4096, 4096),
-        nn.ReLU(inplace=True),
-        nn.Linear(4096, len(classes)),
-    )
-    print(net)
+    # network 2:
+    # net = models.alexnet(pretrained=True)
+    # net.features = nn.Sequential(
+    #     nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
+    #     nn.ReLU(inplace=True),
+    #     nn.MaxPool2d(kernel_size=3, stride=2),
+    #     nn.Conv2d(64, 192, kernel_size=5, padding=2),
+    #     nn.ReLU(inplace=True),
+    #     nn.MaxPool2d(kernel_size=3, stride=2),
+    #     nn.Conv2d(192, 384, kernel_size=3, padding=1),
+    #     nn.ReLU(inplace=True),
+    #     nn.Conv2d(384, 256, kernel_size=3, padding=1),
+    #     nn.ReLU(inplace=True),
+    #     # nn.Conv2d(256, 256, kernel_size=3, padding=1),
+    #     # nn.ReLU(inplace=True),
+    #     nn.MaxPool2d(kernel_size=3, stride=2),
+    # )
+    # net.classifier = nn.Sequential(
+    #     nn.Dropout(),
+    #     nn.Linear(256 * 6 * 6, 4096),
+    #     nn.ReLU(inplace=True),
+    #     nn.Dropout(),
+    #     nn.Linear(4096, 4096),
+    #     nn.ReLU(inplace=True),
+    #     nn.Linear(4096, len(classes)),
+    # )
+    # network 3:
+    net = models.resnet18(pretrained=True)
+    net.fc = nn.Linear(512, len(classes))
 
+    print(net)
     net.to(DEVICE)
     optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)  # learning rate=0.008
     criterion = nn.CrossEntropyLoss()  # loss function