From f28684c463ad8b47b5e2ac9d2db919746258e643 Mon Sep 17 00:00:00 2001 From: "sweep-nightly[bot]" <131841235+sweep-nightly[bot]@users.noreply.github.com> Date: Tue, 24 Oct 2023 21:28:55 +0000 Subject: [PATCH 1/4] feat: Added CNN class and training function in cnn --- src/cnn.py | 41 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) create mode 100644 src/cnn.py diff --git a/src/cnn.py b/src/cnn.py new file mode 100644 index 0000000..a4eac83 --- /dev/null +++ b/src/cnn.py @@ -0,0 +1,41 @@ +import torch +import torch.nn as nn +import torch.optim as optim + +class CNN(nn.Module): + """ + Convolutional Neural Network (CNN) class. + """ + def __init__(self): + super(CNN, self).__init__() + self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1) + self.relu1 = nn.ReLU() + self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) + self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1) + self.relu2 = nn.ReLU() + self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) + self.fc1 = nn.Linear(7 * 7 * 64, 128) + self.fc2 = nn.Linear(128, 10) + + def forward(self, x): + x = self.pool1(self.relu1(self.conv1(x))) + x = self.pool2(self.relu2(self.conv2(x))) + x = x.view(-1, 7 * 7 * 64) + x = self.fc1(x) + x = self.fc2(x) + return x + +def train_cnn(model, dataloader, epochs): + """ + Function to train the CNN model. + """ + optimizer = optim.SGD(model.parameters(), lr=0.01) + criterion = nn.CrossEntropyLoss() + + for epoch in range(epochs): + for images, labels in dataloader: + optimizer.zero_grad() + output = model(images) + loss = criterion(output, labels) + loss.backward() + optimizer.step() From 8f6b33717426abfa3c6ee564e37d8b120eaaf91b Mon Sep 17 00:00:00 2001 From: "sweep-nightly[bot]" <131841235+sweep-nightly[bot]@users.noreply.github.com> Date: Tue, 24 Oct 2023 21:30:21 +0000 Subject: [PATCH 2/4] Sandbox run src/cnn.py --- src/cnn.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/cnn.py b/src/cnn.py index a4eac83..6819040 100644 --- a/src/cnn.py +++ b/src/cnn.py @@ -1,11 +1,12 @@ -import torch import torch.nn as nn import torch.optim as optim + class CNN(nn.Module): """ Convolutional Neural Network (CNN) class. """ + def __init__(self): super(CNN, self).__init__() self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1) @@ -25,6 +26,7 @@ def forward(self, x): x = self.fc2(x) return x + def train_cnn(model, dataloader, epochs): """ Function to train the CNN model. @@ -32,7 +34,7 @@ def train_cnn(model, dataloader, epochs): optimizer = optim.SGD(model.parameters(), lr=0.01) criterion = nn.CrossEntropyLoss() - for epoch in range(epochs): + for _epoch in range(epochs): for images, labels in dataloader: optimizer.zero_grad() output = model(images) From 5db8d2280ccff70073250b51cd1a5a5bb7d0c284 Mon Sep 17 00:00:00 2001 From: "sweep-nightly[bot]" <131841235+sweep-nightly[bot]@users.noreply.github.com> Date: Tue, 24 Oct 2023 21:30:47 +0000 Subject: [PATCH 3/4] feat: Updated src/main.py --- src/main.py | 33 ++++++--------------------------- 1 file changed, 6 insertions(+), 27 deletions(-) diff --git a/src/main.py b/src/main.py index 243a31e..06a965a 100644 --- a/src/main.py +++ b/src/main.py @@ -5,6 +5,7 @@ from torchvision import datasets, transforms from torch.utils.data import DataLoader import numpy as np +from src.cnn import CNN, train_cnn # Import CNN class and train_cnn function from cnn.py # Step 1: Load MNIST Data and Preprocess transform = transforms.Compose([ @@ -16,33 +17,11 @@ 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) +# Create an instance of the CNN class +model = CNN() # Create an instance of the CNN class # Step 3: Train the Model -model = Net() -optimizer = optim.SGD(model.parameters(), lr=0.01) -criterion = nn.NLLLoss() +# Call the train_cnn function, passing the CNN instance, the DataLoader instance, and a number of epochs +train_cnn(model, trainloader, 10) # Train the CNN model with the trainloader data for 10 epochs -# Training loop -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() - -torch.save(model.state_dict(), "mnist_model.pth") \ No newline at end of file +torch.save(model.state_dict(), "mnist_model.pth") # Save the trained model \ No newline at end of file From c59415fad8a51443ffd45b4911738ed4ce77b17a Mon Sep 17 00:00:00 2001 From: "sweep-nightly[bot]" <131841235+sweep-nightly[bot]@users.noreply.github.com> Date: Tue, 24 Oct 2023 21:32:20 +0000 Subject: [PATCH 4/4] Sandbox run src/main.py --- src/main.py | 28 +++++++++++++++------------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/src/main.py b/src/main.py index 06a965a..7bc2620 100644 --- a/src/main.py +++ b/src/main.py @@ -1,19 +1,19 @@ -from PIL import Image +import numpy as np import torch -import torch.nn as nn -import torch.optim as optim -from torchvision import datasets, transforms from torch.utils.data import DataLoader -import numpy as np -from src.cnn import CNN, train_cnn # Import CNN class and train_cnn function from cnn.py +from torchvision import datasets, transforms + +from src.cnn import ( # Import CNN class and train_cnn function from cnn.py + CNN, + train_cnn, +) # Step 1: Load MNIST Data and Preprocess -transform = transforms.Compose([ - transforms.ToTensor(), - transforms.Normalize((0.5,), (0.5,)) -]) +transform = transforms.Compose( + [transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))] +) -trainset = datasets.MNIST('.', download=True, train=True, transform=transform) +trainset = datasets.MNIST(".", download=True, train=True, transform=transform) trainloader = DataLoader(trainset, batch_size=64, shuffle=True) # Step 2: Define the PyTorch Model @@ -22,6 +22,8 @@ # Step 3: Train the Model # Call the train_cnn function, passing the CNN instance, the DataLoader instance, and a number of epochs -train_cnn(model, trainloader, 10) # Train the CNN model with the trainloader data for 10 epochs +train_cnn( + model, trainloader, 10 +) # Train the CNN model with the trainloader data for 10 epochs -torch.save(model.state_dict(), "mnist_model.pth") # Save the trained model \ No newline at end of file +torch.save(model.state_dict(), "mnist_model.pth") # Save the trained model