updated basic tutorials, better comments, code revision, checked it works with latest pytorch version

ML/Pytorch/Basics/pytorch_simple_CNN.py

@@ -8,18 +8,20 @@ check accuracy and more.
 Programmed by Aladdin Persson
 * 2020-04-08: Initial coding
 * 2021-03-24: More detailed comments and small revision of the code
+* 2022-12-19: Small revision of code, checked that it works with latest PyTorch version
 
 """
 
 # Imports
 import torch
-import torchvision # torch package for vision related things
 import torch.nn.functional as F  # Parameterless functions, like (some) activation functions
 import torchvision.datasets as datasets  # Standard datasets
 import torchvision.transforms as transforms  # Transformations we can perform on our dataset for augmentation
 from torch import optim  # For optimizers like SGD, Adam, etc.
 from torch import nn  # All neural network modules
-from torch.utils.data import DataLoader  # Gives easier dataset managment by creating mini batches etc.
+from torch.utils.data import (
+    DataLoader,
+)  # Gives easier dataset managment by creating mini batches etc.
 from tqdm import tqdm  # For nice progress bar!
 
 # Simple CNN
@@ -29,17 +31,17 @@ class CNN(nn.Module):
         self.conv1 = nn.Conv2d(
             in_channels=in_channels,
             out_channels=8,
-            kernel_size=(3, 3),
-            stride=(1, 1),
-            padding=(1, 1),
+            kernel_size=3,
+            stride=1,
+            padding=1,
         )
-        self.pool = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2))
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
         self.conv2 = nn.Conv2d(
             in_channels=8,
             out_channels=16,
-            kernel_size=(3, 3),
-            stride=(1, 1),
-            padding=(1, 1),
+            kernel_size=3,
+            stride=1,
+            padding=1,
         )
         self.fc1 = nn.Linear(16 * 7 * 7, num_classes)
 
@@ -59,13 +61,17 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # Hyperparameters
 in_channels = 1
 num_classes = 10
-learning_rate = 0.001
+learning_rate = 3e-4 # karpathy's constant
 batch_size = 64
 num_epochs = 3
 
 # Load Data
-train_dataset = datasets.MNIST(root="dataset/", train=True, transform=transforms.ToTensor(), download=True)
-test_dataset = datasets.MNIST(root="dataset/", train=False, transform=transforms.ToTensor(), download=True)
+train_dataset = datasets.MNIST(
+    root="dataset/", train=True, transform=transforms.ToTensor(), download=True
+)
+test_dataset = datasets.MNIST(
+    root="dataset/", train=False, transform=transforms.ToTensor(), download=True
+)
 train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
 test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True)
 
@@ -110,10 +116,9 @@ def check_accuracy(loader, model):
             num_correct += (predictions == y).sum()
             num_samples += predictions.size(0)
 
-
     model.train()
-    return num_correct/num_samples
+    return num_correct / num_samples
 
 
 print(f"Accuracy on training set: {check_accuracy(train_loader, model)*100:.2f}")
-print(f"Accuracy on test set: {check_accuracy(test_loader, model)*100:.2f}")
+print(f"Accuracy on test set: {check_accuracy(test_loader, model)*100:.2f}")