[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

computer_vision/mnist_classifier.py

@@ -1,6 +1,6 @@
 """
-This program is a MNIST classifier using AlexNet. It accepts three parameters provided as a command line input. 
-The first two inputs are two digits between 0-9 which are used to train and test the classifier and the third 
+This program is a MNIST classifier using AlexNet. It accepts three parameters provided as a command line input.
+The first two inputs are two digits between 0-9 which are used to train and test the classifier and the third
 parameter controls the number of training epochs.
 Syntax: python program.py <number> <number> <number>
 
@@ -9,7 +9,6 @@
 
 """
 
-
 import sys
 import torch
 import torch.nn as nn
@@ -20,7 +19,7 @@
 import torch.optim as optim
 
 
-class AlexNet(nn.Module):   
+class AlexNet(nn.Module):
     def __init__(self, num=10):
         super(AlexNet, self).__init__()
         self.feature = nn.Sequential(
@@ -36,20 +35,20 @@ def __init__(self, num=10):
             nn.ReLU(inplace=True),
             nn.Conv2d(64, 32, kernel_size=3, padding=1),
             nn.ReLU(inplace=True),
-            nn.MaxPool2d(kernel_size=2, stride=1)         
+            nn.MaxPool2d(kernel_size=2, stride=1),
         )
 
         self.classifier = nn.Sequential(
             # Define classifier here...
             nn.Dropout(),
-            nn.Linear(32*12*12, 2048),
+            nn.Linear(32 * 12 * 12, 2048),
             nn.ReLU(inplace=True),
             nn.Dropout(),
             nn.Linear(2048, 1024),
             nn.ReLU(inplace=True),
-            nn.Linear(1024, 10)
+            nn.Linear(1024, 10),
         )
-    
+
     def forward(self, x):
         # define forward network 'x' that combines feature extractor and classifier
         x = self.feature(x)
@@ -63,26 +62,27 @@ def load_subset(full_train_set, full_test_set, label_one, label_two):
     train_set = []
     data_lim = 20000
     for data in full_train_set:
-        if data_lim>0:
-            data_lim-=1
-            if data[1]==label_one or data[1]==label_two:
+        if data_lim > 0:
+            data_lim -= 1
+            if data[1] == label_one or data[1] == label_two:
                 train_set.append(data)
         else:
             break
 
     test_set = []
     data_lim = 1000
     for data in full_test_set:
-        if data_lim>0:
-            data_lim-=1
-            if data[1]==label_one or data[1]==label_two:
+        if data_lim > 0:
+            data_lim -= 1
+            if data[1] == label_one or data[1] == label_two:
                 test_set.append(data)
         else:
             break
 
     return train_set, test_set
 
-def train(model,optimizer,train_loader,epoch):
+
+def train(model, optimizer, train_loader, epoch):
     model.train()
     for batch_idx, (data, target) in enumerate(train_loader):
         if torch.cuda.is_available():
@@ -94,7 +94,8 @@ def train(model,optimizer,train_loader,epoch):
         loss.backward()
         optimizer.step()
 
-def test(model,test_loader):
+
+def test(model, test_loader):
     model.eval()
     test_loss = 0
     correct = 0
@@ -104,65 +105,76 @@ def test(model,test_loader):
         with torch.no_grad():
             data, target = Variable(data), Variable(target)
         output = model(data)
-        test_loss += F.cross_entropy(output, target, reduction='sum').item()#size_average=False
-        pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability
+        test_loss += F.cross_entropy(
+            output, target, reduction="sum"
+        ).item()  # size_average=False
+        pred = output.data.max(1, keepdim=True)[
+            1
+        ]  # get the index of the max log-probability
         correct += pred.eq(target.data.view_as(pred)).long().cpu().sum()
-    
+
     test_loss /= len(test_loader.dataset)
-    acc=100. * float(correct.to(torch.device('cpu')).numpy())
-    test_accuracy = (acc / len(test_loader.dataset))
+    acc = 100.0 * float(correct.to(torch.device("cpu")).numpy())
+    test_accuracy = acc / len(test_loader.dataset)
     return test_accuracy
 
 
 """ Start to call """
 
-if __name__ == '__main__':
- 
+if __name__ == "__main__":
     if len(sys.argv) == 3:
         print("Usage: python assignment.py <number> <number>")
         sys.exit(1)
 
     input_data_one = sys.argv[1].strip()
     input_data_two = sys.argv[2].strip()
     epochs = sys.argv[3].strip()
-    
+
     """  Call to function that will perform the computation. """
     if input_data_one.isdigit() and input_data_two.isdigit() and epochs.isdigit():
-
         label_one = int(input_data_one)
         label_two = int(input_data_two)
         epochs = int(epochs)
-        
-        if label_one!=label_two and 0<=label_one<=9 and 0<=label_two<=9:
+
+        if label_one != label_two and 0 <= label_one <= 9 and 0 <= label_two <= 9:
             torch.manual_seed(42)
             # Load MNIST dataset
-            trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
-            full_train_set = dset.MNIST(root='./data', train=True, transform=trans, download=True)
-            full_test_set = dset.MNIST(root='./data', train=False, transform=trans)
+            trans = transforms.Compose(
+                [transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))]
+            )
+            full_train_set = dset.MNIST(
+                root="./data", train=True, transform=trans, download=True
+            )
+            full_test_set = dset.MNIST(root="./data", train=False, transform=trans)
             batch_size = 16
             # Get final train and test sets
-            train_set, test_set = load_subset(full_train_set,full_test_set,label_one,label_two)
-            
-            train_loader = torch.utils.data.DataLoader(dataset=train_set,batch_size=batch_size,shuffle=False)
-            test_loader = torch.utils.data.DataLoader(dataset=test_set,batch_size=batch_size,shuffle=False)
+            train_set, test_set = load_subset(
+                full_train_set, full_test_set, label_one, label_two
+            )
+
+            train_loader = torch.utils.data.DataLoader(
+                dataset=train_set, batch_size=batch_size, shuffle=False
+            )
+            test_loader = torch.utils.data.DataLoader(
+                dataset=test_set, batch_size=batch_size, shuffle=False
+            )
 
             model = AlexNet()
             if torch.cuda.is_available():
                 model.cuda()
-                    
+
             optimizer = optim.SGD(model.parameters(), lr=0.01)
-            
-            for epoch in range(1, epochs+1):
-                train(model,optimizer,train_loader,epoch)
-                accuracy = test(model,test_loader)
-
-            print(round(accuracy,2))
-            
-            
+
+            for epoch in range(1, epochs + 1):
+                train(model, optimizer, train_loader, epoch)
+                accuracy = test(model, test_loader)
+
+            print(round(accuracy, 2))
+
         else:
-           print("Invalid input")
+            print("Invalid input")
     else:
         print("Invalid input")
- 
-    
+
+
 """ End to call """