feat: Updated src/main.py

src/main.py

@@ -5,6 +5,7 @@
 from torchvision import datasets, transforms
 from torch.utils.data import DataLoader
 import numpy as np
+from cnn import CNN
 
 # Step 1: Load MNIST Data and Preprocess
 transform = transforms.Compose([
@@ -15,34 +16,15 @@
 trainset = datasets.MNIST('.', download=True, train=True, transform=transform)
 trainloader = DataLoader(trainset, batch_size=64, shuffle=True)
 
-# Step 2: Define the PyTorch Model
-class Net(nn.Module):
- def __init__(self):
- super().__init__()
- self.fc1 = nn.Linear(28 * 28, 128)
- self.fc2 = nn.Linear(128, 64)
- self.fc3 = nn.Linear(64, 10)
-
- def forward(self, x):
- x = x.view(-1, 28 * 28)
- x = nn.functional.relu(self.fc1(x))
- x = nn.functional.relu(self.fc2(x))
- x = self.fc3(x)
- return nn.functional.log_softmax(x, dim=1)
+# Step 2: Instantiate the CNN Model
+model = CNN()
 
-# Step 3: Train the Model
-model = Net()
+# Step 3: Define the Optimizer
 optimizer = optim.SGD(model.parameters(), lr=0.01)
-criterion = nn.NLLLoss()
 
-# Training loop
+# Step 4: Train the Model using the CNN's train method
 epochs = 3
 for epoch in range(epochs):
- for images, labels in trainloader:
- optimizer.zero_grad()
- output = model(images)
- loss = criterion(output, labels)
- loss.backward()
- optimizer.step()
+ model.train(trainloader, optimizer)
 
 torch.save(model.state_dict(), "mnist_model.pth")