From dff26a609880dd195da22b6ce7815cde9857ceae Mon Sep 17 00:00:00 2001 From: "sweep-nightly[bot]" <131841235+sweep-nightly[bot]@users.noreply.github.com> Date: Thu, 12 Oct 2023 23:03:19 +0000 Subject: [PATCH 1/4] feat: Updated src/main.py --- src/main.py | 99 +++++++++++++++++++++++++++++++---------------------- 1 file changed, 59 insertions(+), 40 deletions(-) diff --git a/src/main.py b/src/main.py index 243a31e..121902e 100644 --- a/src/main.py +++ b/src/main.py @@ -6,43 +6,62 @@ from torch.utils.data import DataLoader import numpy as np -# Step 1: Load MNIST Data and Preprocess -transform = transforms.Compose([ - transforms.ToTensor(), - transforms.Normalize((0.5,), (0.5,)) -]) - -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 3: Train the Model -model = Net() -optimizer = optim.SGD(model.parameters(), lr=0.01) -criterion = nn.NLLLoss() - -# 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 +class ModelTrainer: + def __init__(self, epochs=3, lr=0.01, batch_size=64): + self.epochs = epochs + self.lr = lr + self.batch_size = batch_size + self.model = None + self.optimizer = None + self.criterion = nn.NLLLoss() + self.trainloader = None + + def load_data(self): + transform = transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize((0.5,), (0.5,)) + ]) + + trainset = datasets.MNIST('.', download=True, train=True, transform=transform) + self.trainloader = DataLoader(trainset, batch_size=self.batch_size, shuffle=True) + + def define_model(self): + 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) + + self.model = Net() + self.optimizer = optim.SGD(self.model.parameters(), lr=self.lr) + + def train_model(self): + for epoch in range(self.epochs): + for images, labels in self.trainloader: + self.optimizer.zero_grad() + output = self.model(images) + loss = self.criterion(output, labels) + loss.backward() + self.optimizer.step() + + def evaluate(self): + # TODO: Implement model evaluation + pass + + def run(self): + self.load_data() + self.define_model() + self.train_model() + self.evaluate() + torch.save(self.model.state_dict(), "mnist_model.pth") + +trainer = ModelTrainer() +trainer.run() \ No newline at end of file From cc5ebf2966807c0e73b1d7e595749fec12436c7e Mon Sep 17 00:00:00 2001 From: "sweep-nightly[bot]" <131841235+sweep-nightly[bot]@users.noreply.github.com> Date: Thu, 12 Oct 2023 23:04:33 +0000 Subject: [PATCH 2/4] feat: Updated src/api.py --- src/api.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/api.py b/src/api.py index 36c257a..0bf2b8b 100644 --- a/src/api.py +++ b/src/api.py @@ -2,12 +2,11 @@ from PIL import Image import torch from torchvision import transforms -from main import Net # Importing Net class from main.py +from main import ModelTrainer # Importing ModelTrainer class from main.py -# Load the model -model = Net() -model.load_state_dict(torch.load("mnist_model.pth")) -model.eval() +# Create an instance of ModelTrainer and load the model +trainer = ModelTrainer() +model = trainer.load_model("mnist_model.pth") # Transform used for preprocessing the image transform = transforms.Compose([ From 59cf840852afa26842e33cbc57925724233109ff Mon Sep 17 00:00:00 2001 From: "sweep-nightly[bot]" <131841235+sweep-nightly[bot]@users.noreply.github.com> Date: Thu, 12 Oct 2023 23:06:29 +0000 Subject: [PATCH 3/4] feat: Updated README.md --- README.md | 34 +++++++++++++++++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index ea3afcc..d81eb3f 100644 --- a/README.md +++ b/README.md @@ -1 +1,33 @@ -# evals \ No newline at end of file +# evals + +This project contains a PyTorch implementation for training and evaluating a simple neural network on the MNIST dataset. + +## ModelTrainer class + +The `ModelTrainer` class is used to train and evaluate the model. Here is how you can use it: + +```python +from main import ModelTrainer + +# Create an instance of ModelTrainer +trainer = ModelTrainer() + +# Run the training and evaluation process +trainer.run() +``` + +The `ModelTrainer` class has the following methods: + +- `load_data()`: Loads the MNIST dataset and applies the necessary transformations. +- `define_model()`: Defines the model architecture. +- `train_model()`: Trains the model on the MNIST dataset. +- `evaluate()`: Evaluates the model performance on the test dataset. +- `run()`: Runs the entire process (data loading, model definition, training, and evaluation). + +## Saving and Loading the Model + +After training, the model's parameters are saved to a file named "mnist_model.pth". You can load the model parameters using the `torch.load()` function: + +```python +model.load_state_dict(torch.load("mnist_model.pth")) +``` \ No newline at end of file From 62b5b7c62fe1e4a1b37d0b3cc0346dbd505ce98e Mon Sep 17 00:00:00 2001 From: "sweep-nightly[bot]" <131841235+sweep-nightly[bot]@users.noreply.github.com> Date: Thu, 12 Oct 2023 23:08:54 +0000 Subject: [PATCH 4/4] feat: Updated src/main.py --- src/main.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/src/main.py b/src/main.py index 121902e..06bee82 100644 --- a/src/main.py +++ b/src/main.py @@ -53,8 +53,26 @@ def train_model(self): self.optimizer.step() def evaluate(self): - # TODO: Implement model evaluation - pass + transform = transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize((0.5,), (0.5,)) + ]) + + testset = datasets.MNIST('.', download=True, train=False, transform=transform) + testloader = DataLoader(testset, batch_size=self.batch_size, shuffle=True) + + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in testloader: + output = self.model(images) + _, predicted = torch.max(output.data, 1) + total += labels.size(0) + correct += (predicted == labels).sum().item() + + accuracy = 100 * correct / total + print('Accuracy of the model on the test images: %d%%' % accuracy) def run(self): self.load_data()