Certainly! Here is the complete code, including the README content, training, evaluation, and saving/loading the model, all in a single script:
# README content as a comment
"""
# MNIST Handwritten Digit Classification
This project involves training a neural network to classify handwritten digits from the MNIST dataset. The model is implemented using PyTorch and achieves high accuracy on the test set.
## Project Overview
- **Dataset**: MNIST (Modified National Institute of Standards and Technology) dataset, which contains 60,000 training images and 10,000 testing images of handwritten digits (0-9).
- **Model Architecture**: A simple Multi-Layer Perceptron (MLP) neural network.
- **Framework**: PyTorch.
## Preprocessing
The preprocessing steps include:
1. **Converting to Tensors**: Images are converted to PyTorch tensors using `transforms.ToTensor()`.
2. **Normalization**: Pixel values are normalized to the range [-1, 1] using `transforms.Normalize((0.5,), (0.5,))`.
3. **Flattening**: Images are flattened from 28x28 pixels to a single 784-dimensional vector before being fed into the network.
## Model Architecture
The MLP model consists of the following layers:
- Input layer: 784 neurons (for 28x28 pixel images)
- Hidden layer 1: 120 neurons with ReLU activation
- Hidden layer 2: 84 neurons with ReLU activation
- Output layer: 10 neurons with log-softmax activation (for 10 classes)
## Training
The model is trained using the following settings:
- **Loss Function**: CrossEntropyLoss
- **Optimizer**: Stochastic Gradient Descent (SGD) with learning rate 0.005 and momentum 0.9
- **Number of Epochs**: 10
## Evaluation
The model is evaluated on the test set, and the accuracy is computed by comparing the predicted labels with the true labels.
## Usage
1. **Clone the repository**:
```sh
git clone https://github.com/yourusername/mnist-classification.git
cd mnist-classification-
Install dependencies: Make sure you have PyTorch installed. You can install the other dependencies using:
pip install -r requirements.txt
-
Run the training script:
python train.py
-
Evaluate the model:
python evaluate.py
The model achieves an accuracy of over 97% on the test set after 10 epochs of training.
This project is licensed under the MIT License.
- The MNIST dataset is provided by Yann LeCun, Corinna Cortes, and Chris Burges.
- PyTorch is an open-source machine learning library developed by Facebook's AI Research lab. """
import torch from torch import nn, optim from torch.utils.data import DataLoader from torchvision import datasets, transforms import matplotlib.pyplot as plt
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
training_data = datasets.MNIST(root="data", train=True, download=True, transform=transform) train_loader = DataLoader(training_data, batch_size=32, shuffle=True) testing_data = datasets.MNIST(root="data", train=False, download=True, transform=transform) test_loader = DataLoader(testing_data, batch_size=32)
class MNIST_MLP(nn.Module): def init(self): super(MNIST_MLP, self).init() self.fc1 = nn.Linear(28 * 28, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) self.activation = nn.ReLU()
def forward(self, x):
x = torch.flatten(x, 1)
x = self.activation(self.fc1(x))
x = self.activation(self.fc2(x))
x = self.fc3(x)
return x
mlp = MNIST_MLP() if torch.cuda.is_available(): mlp.cuda()
def train_network_classification(net, train_loader, test_loader): num_epochs = 10 criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(net.parameters(), lr=0.005, momentum=0.9) train_loss_history = [] val_loss_history = []
for epoch in range(num_epochs):
net.train()
train_loss = 0.0
train_correct = 0
for inputs, labels in train_loader:
if torch.cuda.is_available():
inputs, labels = inputs.cuda(), labels.cuda()
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
_, preds = torch.max(outputs.data, 1)
train_correct += (preds == labels).sum().item()
train_loss += loss.item()
print(f'Epoch {epoch + 1} training accuracy: {train_correct / len(training_data) * 100:.2f}% training loss: {train_loss / len(train_loader):.5f}')
train_loss_history.append(train_loss / len(train_loader))
val_loss = 0.0
val_correct = 0
net.eval()
with torch.no_grad():
for inputs, labels in test_loader:
if torch.cuda.is_available():
inputs, labels = inputs.cuda(), labels.cuda()
outputs = net(inputs)
loss = criterion(outputs, labels)
_, preds = torch.max(outputs.data, 1)
val_correct += (preds == labels).sum().item()
val_loss += loss.item()
print(f'Epoch {epoch + 1} validation accuracy: {val_correct / len(testing_data) * 100:.2f}% validation loss: {val_loss / len(test_loader):.5f}')
val_loss_history.append(val_loss / len(test_loader))
plt.plot(train_loss_history, label="Training Loss")
plt.plot(val_loss_history, label="Validation Loss")
plt.legend()
plt.show()
# Save the model
torch.save(net.state_dict(), 'mnist_mlp.pth')
train_network_classification(mlp, train_loader, test_loader)
loaded_model = MNIST_MLP() loaded_model.load_state_dict(torch.load('mnist_mlp.pth')) if torch.cuda.is_available(): loaded_model.cuda() loaded_model.eval()
def evaluate_model(net, test_loader): net.eval() correct = 0 total = 0 with torch.no_grad(): for data in test_loader: images, labels = data if torch.cuda.is_available(): images, labels = images.cuda(), labels.cuda() outputs = net(images) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() accuracy = (correct / total) * 100 return accuracy
test_accuracy = evaluate_model(loaded_model, test_loader) print(f'Test Accuracy: {test_accuracy:.2f}%') which is 98% accuracy