Refactor training loop from script to class

src/api.py

@@ -2,7 +2,7 @@
 from PIL import Image
 import torch
 from torchvision import transforms
-from main import Net # Importing Net class from main.py
+from main import Net, TrainModel # Importing Net and TrainModel classes from main.py
 
 # Load the model
 model = Net()
@@ -14,6 +14,9 @@
  transforms.ToTensor(),
  transforms.Normalize((0.5,), (0.5,))
 ])
+# Create an instance of TrainModel
+train_model = TrainModel(model, criterion, optimizer, trainloader)
+train_model.train(3)
 
 app = FastAPI()
 
@@ -23,6 +26,6 @@ async def predict(file: UploadFile = File(...)):
  image = transform(image)
  image = image.unsqueeze(0) # Add batch dimension
  with torch.no_grad():
- output = model(image)
+ output = train_model.model(image)
  _, predicted = torch.max(output.data, 1)
  return {"prediction": int(predicted[0])}

src/main.py

@@ -15,6 +15,23 @@
 trainset = datasets.MNIST('.', download=True, train=True, transform=transform)
 trainloader = DataLoader(trainset, batch_size=64, shuffle=True)
 
+# Step 4: Define the TrainModel Class
+class TrainModel:
+ def __init__(self, model, criterion, optimizer, dataloader):
+ self.model = model
+ self.criterion = criterion
+ self.optimizer = optimizer
+ self.dataloader = dataloader
+
+ def train(self, epochs):
+ for epoch in range(epochs):
+ for images, labels in self.dataloader:
+ self.optimizer.zero_grad()
+ output = self.model(images)
+ loss = self.criterion(output, labels)
+ loss.backward()
+ self.optimizer.step()
+
 # Step 2: Define the PyTorch Model
 class Net(nn.Module):
  def __init__(self):
@@ -35,14 +52,8 @@ def forward(self, x):
 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()
+# Create an instance of TrainModel and train
+train_model = TrainModel(model, criterion, optimizer, trainloader)
+train_model.train(3)
 
 torch.save(model.state_dict(), "mnist_model.pth")