Rename pages/1_MNIST_With_Images.py to pages/1_Simple_CNN.py

pages/{1_MNIST_With_Images.py → 1_Simple_CNN.py}

@@ -5,49 +5,51 @@ import torch.optim as optim
 import torchvision
 import torchvision.transforms as transforms
 import matplotlib.pyplot as plt
+import numpy as np
 
-# Define the neural network
-class Net(nn.Module):
+# Define the CNN
+class SimpleCNN(nn.Module):
     def __init__(self):
-        super(Net, self).__init__()
-        self.fc1 = nn.Linear(28 * 28, 128)
-        self.fc2 = nn.Linear(128, 64)
-        self.fc3 = nn.Linear(64, 10)
+        super(SimpleCNN, self).__init__()
+        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
+        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.fc1 = nn.Linear(32 * 8 * 8, 128)
+        self.fc2 = nn.Linear(128, 10)
     
     def forward(self, x):
-        x = x.view(-1, 28 * 28)
+        x = self.pool(torch.relu(self.conv1(x)))
+        x = self.pool(torch.relu(self.conv2(x)))
+        x = x.view(-1, 32 * 8 * 8)
         x = torch.relu(self.fc1(x))
-        x = torch.relu(self.fc2(x))
-        x = self.fc3(x)
+        x = self.fc2(x)
         return x
 
 # Function to train the model
 def train_model(num_epochs):
-    # Define transformations
     transform = transforms.Compose([
         transforms.ToTensor(),
-        transforms.Normalize((0.5,), (0.5,))
+        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
     ])
 
-    # Load datasets
-    trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
+    trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
     trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
 
-    testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
+    testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
     testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)
 
-    # Display sample images
-    display_sample_images(trainloader)
+    CIFAR10_CLASSES = [
+        'plane', 'car', 'bird', 'cat', 'deer',
+        'dog', 'frog', 'horse', 'ship', 'truck'
+    ]
 
-    # Initialize the network, loss function, and optimizer
-    net = Net()
+    net = SimpleCNN()
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)
 
-    # Track loss over epochs
     loss_values = []
+    st.write("Training the model...")
 
-    # Training loop
     for epoch in range(num_epochs):
         running_loss = 0.0
         for i, data in enumerate(trainloader, 0):
@@ -58,11 +60,8 @@ def train_model(num_epochs):
             loss.backward()
             optimizer.step()
             running_loss += loss.item()
-        
-        # Append average loss for this epoch
         loss_values.append(running_loss / len(trainloader))
         st.write(f'Epoch {epoch + 1}, Loss: {running_loss / len(trainloader):.3f}')
-
     st.write('Finished Training')
 
     # Plot the loss values
@@ -73,36 +72,39 @@ def train_model(num_epochs):
     plt.ylabel('Loss')
     st.pyplot(plt)
 
-    # Evaluate the network on the test data
     correct = 0
     total = 0
     with torch.no_grad():
         for data in testloader:
             images, labels = data
             outputs = net(images)
-            _, predicted = torch.max(outputs.data, 1)
+            _, predicted = torch.max(outputs, 1)
             total += labels.size(0)
             correct += (predicted == labels).sum().item()
 
     st.write(f'Accuracy of the network on the 10000 test images: {100 * correct / total}%')
 
-# Function to display sample images
-def display_sample_images(loader):
-    dataiter = iter(loader)
+    # Visualize some test images and their predictions
+    def imshow(img):
+        img = img / 2 + 0.5  # Unnormalize
+        npimg = img.numpy()
+        plt.imshow(np.transpose(npimg, (1, 2, 0)))
+        plt.show()
+
+    dataiter = iter(testloader)
     images, labels = dataiter.next()
-    images = images[:5]
-    labels = labels[:5]
 
-    fig, axes = plt.subplots(1, 5, figsize=(15, 3))
-    for i in range(5):
-        ax = axes[i]
-        ax.imshow(images[i].numpy().squeeze(), cmap='gray')
-        ax.set_title(f'Label: {labels[i].item()}')
-        ax.axis('off')
-    st.pyplot(fig)
+    imshow(torchvision.utils.make_grid(images))
+
+    outputs = net(images)
+    _, predicted = torch.max(outputs, 1)
+
+    st.write('Predicted: ', ' '.join(f'{CIFAR10_CLASSES[predicted[j]]:5s}' for j in range(8)))
+    st.write('Actual:    ', ' '.join(f'{CIFAR10_CLASSES[labels[j]]:5s}' for j in range(8)))
+    st.pyplot()
 
 # Streamlit interface
-st.title('MNIST Digit Classification with PyTorch')
+st.title('CIFAR-10 Classification with PyTorch')
 num_epochs = st.number_input('Enter number of epochs:', min_value=1, max_value=100, value=10)
 if st.button('Run'):
     train_model(num_epochs)