Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor training loop from script to class #70

Closed
wants to merge 4 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 8 additions & 14 deletions src/api.py
Original file line number Diff line number Diff line change
@@ -1,26 +1,20 @@
from fastapi import FastAPI, UploadFile, File
from PIL import Image
import torch
from torchvision import transforms
from main import Net # Importing Net class from main.py
from fastapi import FastAPI, File, UploadFile
from PIL import Image

# Load the model
model = Net()
model.load_state_dict(torch.load("mnist_model.pth"))
model.eval()
from main import MNISTTrainer # Importing MNISTTrainer class from main.py

# 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()


@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)
Expand Down
95 changes: 54 additions & 41 deletions src/main.py
Original file line number Diff line number Diff line change
@@ -1,48 +1,61 @@
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

# Step 1: Load MNIST Data and Preprocess
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
from torchvision import datasets, transforms

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):
class MNISTTrainer:
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")
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
)
trainloader = DataLoader(trainset, batch_size=64, shuffle=True)
return trainloader

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)

model = Net()
return model

def train_model(self, trainloader, model, epochs, lr):
optimizer = optim.SGD(model.parameters(), lr=lr)
criterion = nn.NLLLoss()
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, epochs=3, lr=0.01, save_path="mnist_model.pth"):
trainloader = self.load_data()
model = self.define_model()
model = self.train_model(trainloader, model, epochs, lr)
torch.save(model.state_dict(), save_path)
return model


trainer = MNISTTrainer()
trainer.run()
Loading