-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
4d62685
commit 4f2f24f
Showing
1 changed file
with
64 additions
and
47 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,48 +1,65 @@ | ||
from PIL import Image | ||
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,)) | ||
]) | ||
|
||
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,)) | ||
]) | ||
self.optimizer = None | ||
self.criterion = nn.NLLLoss() | ||
self.epochs = 3 | ||
|
||
def load_data(self): | ||
"""Load and preprocess MNIST data.""" | ||
trainset = datasets.MNIST('.', download=True, train=True, transform=self.transform) | ||
trainloader = DataLoader(trainset, batch_size=64, shuffle=True) | ||
return trainloader | ||
|
||
class Net(nn.Module): | ||
"""Define the PyTorch Model.""" | ||
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): | ||
"""Define the model.""" | ||
model = self.Net() | ||
self.optimizer = optim.SGD(model.parameters(), lr=0.01) | ||
return model | ||
|
||
def train_model(self, model, trainloader): | ||
"""Train the model.""" | ||
for epoch in range(self.epochs): | ||
for images, labels in trainloader: | ||
self.optimizer.zero_grad() | ||
output = model(images) | ||
loss = self.criterion(output, labels) | ||
loss.backward() | ||
self.optimizer.step() | ||
|
||
def save_model(self, model): | ||
"""Save the trained model.""" | ||
torch.save(model.state_dict(), "mnist_model.pth") | ||
|
||
# Create an instance of MNISTTrainer | ||
trainer = MNISTTrainer() | ||
|
||
# Load the data | ||
trainloader = trainer.load_data() | ||
|
||
# Define the model | ||
model = trainer.define_model() | ||
|
||
# Train the model | ||
trainer.train_model(model, trainloader) | ||
|
||
# Save the model | ||
trainer.save_model(model) |