From a1f27c249e42652e9e9e14c8779d2886ff3161ff Mon Sep 17 00:00:00 2001 From: "sweep-nightly[bot]" <131841235+sweep-nightly[bot]@users.noreply.github.com> Date: Mon, 16 Oct 2023 05:37:23 +0000 Subject: [PATCH 1/4] feat: Updated src/main.py --- src/main.py | 79 +++++++++++++++++++++++++++++++---------------------- 1 file changed, 47 insertions(+), 32 deletions(-) diff --git a/src/main.py b/src/main.py index 243a31e..46d0741 100644 --- a/src/main.py +++ b/src/main.py @@ -6,43 +6,58 @@ 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,)) -]) +class MNISTTrainer: + def __init__(self): + self.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) + def load_data(self): + trainset = datasets.MNIST('.', download=True, train=True, transform=self.transform) + trainloader = DataLoader(trainset, batch_size=64, shuffle=True) + return trainloader -# 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) + 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) + return Net() # 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 + def train_model(self, trainloader, model): + optimizer = optim.SGD(model.parameters(), lr=0.01) + criterion = nn.NLLLoss() + 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() + return model + + def run(self): + trainloader = self.load_data() + model = self.define_model() + model = self.train_model(trainloader, model) + torch.save(model.state_dict(), "mnist_model.pth") + return model + +trainer = MNISTTrainer() +trainer.run() \ No newline at end of file From 8dde0b53420df63d993a6cebc3436b33ce3dbdf9 Mon Sep 17 00:00:00 2001 From: "sweep-nightly[bot]" <131841235+sweep-nightly[bot]@users.noreply.github.com> Date: Mon, 16 Oct 2023 05:41:59 +0000 Subject: [PATCH 2/4] feat: Updated src/main.py --- src/main.py | 38 ++++++++++++++++++-------------------- 1 file changed, 18 insertions(+), 20 deletions(-) diff --git a/src/main.py b/src/main.py index 46d0741..3c3b22f 100644 --- a/src/main.py +++ b/src/main.py @@ -1,20 +1,21 @@ -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 torchvision import datasets, transforms + class MNISTTrainer: def __init__(self): - self.transform = transforms.Compose([ - transforms.ToTensor(), - transforms.Normalize((0.5,), (0.5,)) - ]) + self.transform = transforms.Compose( + [transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))] + ) def load_data(self): - trainset = datasets.MNIST('.', download=True, train=True, transform=self.transform) + trainset = datasets.MNIST( + ".", download=True, train=True, transform=self.transform + ) trainloader = DataLoader(trainset, batch_size=64, shuffle=True) return trainloader @@ -32,17 +33,13 @@ def forward(self, x): x = nn.functional.relu(self.fc2(x)) x = self.fc3(x) return nn.functional.log_softmax(x, dim=1) - return Net() -# Step 3: Train the Model -model = Net() -optimizer = optim.SGD(model.parameters(), lr=0.01) -criterion = nn.NLLLoss() + model = Net() + return model - def train_model(self, trainloader, model): - optimizer = optim.SGD(model.parameters(), lr=0.01) + def train_model(self, trainloader, model, epochs, lr): + optimizer = optim.SGD(model.parameters(), lr=lr) criterion = nn.NLLLoss() - epochs = 3 for epoch in range(epochs): for images, labels in trainloader: optimizer.zero_grad() @@ -52,12 +49,13 @@ def train_model(self, trainloader, model): optimizer.step() return model - def run(self): + def run(self, epochs=3, lr=0.01, save_path="mnist_model.pth"): trainloader = self.load_data() model = self.define_model() - model = self.train_model(trainloader, model) - torch.save(model.state_dict(), "mnist_model.pth") + model = self.train_model(trainloader, model, epochs, lr) + torch.save(model.state_dict(), save_path) return model + trainer = MNISTTrainer() -trainer.run() \ No newline at end of file +trainer.run() From ab8aa3031f966d4cbe210cb13780b96846861805 Mon Sep 17 00:00:00 2001 From: "sweep-nightly[bot]" <131841235+sweep-nightly[bot]@users.noreply.github.com> Date: Mon, 16 Oct 2023 05:43:53 +0000 Subject: [PATCH 3/4] feat: Updated src/api.py --- src/api.py | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/src/api.py b/src/api.py index 36c257a..1c7c30f 100644 --- a/src/api.py +++ b/src/api.py @@ -2,18 +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 MNISTTrainer # Importing MNISTTrainer class from main.py -# Load the model -model = Net() -model.load_state_dict(torch.load("mnist_model.pth")) -model.eval() - -# Transform used for preprocessing the image -transform = transforms.Compose([ - transforms.ToTensor(), - transforms.Normalize((0.5,), (0.5,)) -]) +# Create an instance of MNISTTrainer and run the training +trainer = MNISTTrainer() +model = trainer.run() app = FastAPI() From a676bbe7a7dc94dff380583468a8465134b284f5 Mon Sep 17 00:00:00 2001 From: "sweep-nightly[bot]" <131841235+sweep-nightly[bot]@users.noreply.github.com> Date: Mon, 16 Oct 2023 05:49:18 +0000 Subject: [PATCH 4/4] feat: Updated src/api.py --- src/api.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/api.py b/src/api.py index 1c7c30f..fae9a17 100644 --- a/src/api.py +++ b/src/api.py @@ -1,7 +1,7 @@ -from fastapi import FastAPI, UploadFile, File -from PIL import Image import torch -from torchvision import transforms +from fastapi import FastAPI, File, UploadFile +from PIL import Image + from main import MNISTTrainer # Importing MNISTTrainer class from main.py # Create an instance of MNISTTrainer and run the training @@ -10,10 +10,11 @@ app = FastAPI() + @app.post("/predict/") async def predict(file: UploadFile = File(...)): image = Image.open(file.file).convert("L") - image = transform(image) + image = trainer.preprocess(image) image = image.unsqueeze(0) # Add batch dimension with torch.no_grad(): output = model(image)