Hey community! If you're just starting out with artificial intelligence (AI) and want to develop your own AI model, you've come to the right place! Today, I'll walk you through step-by-step how to develop and train a basic AI model with PyTorch, a favorite for creating intelligent programs. No worries if you're a beginner—basically, I'll use easy-to-understand language and make it engaging and simple to follow. By the end, you’ll have your own model that can recognize handwritten numbers, and you’ll feel proud of what you’ve achieved!
What Are We Doing? Think about how you could teach a computer to read your handwriting, such as the digits you scrawl onto a piece of paper. We're going to work with a dataset called MNIST, which contains many images of handwritten digits (0 to 9). Our AI program will be able to study these images and make an educated guess about what digit is scribbled down. This is an excellent first project because it's easy but gets you started on the foundation of AI.
Step 1: Get Your Tools Ready
1. Install Python: - Python is like the language the computer understands. Download the latest version (like 3.13 or 3.12) from python.org. When installing, check the box that says “Add Python to PATH” so it works easily.
2. Set Up VS Code: - Download VS Code from code.visualstudio.com and install it. This is where you’ll write your code. Once installed, open it and install the “Python” extension by Microsoft (search for it in the Extensions tab on the left).
3. Create a Project Folder: - Make a new folder on your desktop (e.g., “AI-Model-Train-PyTorch”) to keep all your files together. Open this folder in VS Code by going to File > Open Folder.
4. Install PyTorch:- PyTorch is the magic tool for building AI models.
Open a terminal in VS Code (Terminal > New Terminal) and type:
python -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
If it doesn’t work, try an older Python version (like 3.12) and use:
python -m pip install torch torchvision
python -c "import torch; print(torch.__version__)"
You should see a version number like “2.5.0.dev20250924.”
Step 2: Write Your First AI Code Now, let’s create a file to build our model. This code will teach the computer to recognize handwritten digits.
1. Create the File: - In VS Code, right-click in the Explorer (left side), select New File, and name it mnist_classifier.py.
2. Copy This Simple Code:
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
# 1. Use the computer’s brain (CPU for now)
device = torch.device('cpu')
# 2. Get the MNIST dataset (pictures of digits)
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
# 3. Build a simple model (like a smart guessing machine)
class SimpleNet(nn.Module):
def __init__(self):
super(SimpleNet, self).__init__()
self.flatten = nn.Flatten()
self.fc1 = nn.Linear(28 * 28, 100)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(100, 10)
def forward(self, x):
x = self.flatten(x)
x = self.relu(self.fc1(x))
x = self.fc2(x)
return x
model = SimpleNet().to(device)
# 4. Set up a teacher and a coach
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 5. Train the model (teach it with examples)
num_epochs = 3
for epoch in range(num_epochs):
model.train()
for images, labels in train_loader:
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
print(f'Epoch {epoch+1}, Loss: {loss.item():.4f}')
# 6. Test the model (see how good it is)
model.eval()
correct = 0
total = 0
with torch.no_grad():
for images, labels in test_loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
accuracy = 100 * correct / total
print(f'Test Accuracy: {accuracy:.2f}%')
# 7. Show some pictures (optional)
import matplotlib.pyplot as plt
images, labels = next(iter(test_loader))
images, labels = images.to(device), labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs, 1)
fig, axes = plt.subplots(1, 5, figsize=(10, 2))
for i in range(5):
axes[i].imshow(images[i].cpu().squeeze(), cmap='gray')
axes[i].set_title(f'True: {labels[i].item()}\nPred: {predicted[i].item()}')
axes[i].axis('off')
plt.show()
3. Save the File: - Click Ctrl+S to save.
Step 3: Run Your AI Model
1. Run the Code: - In VS Code, click the green play button (top-right) or type in the terminal:
python mnist_classifier.py
The first time, it will download the MNIST dataset (about 11MB) to a data folder. This might take a minute.
2. What Happens?:
- The model trains for 3 “epochs” (like 3 rounds of learning).
- You’ll see messages like Epoch 1, Loss: 0.3121 and finally Test Accuracy: 95.37%.
- A window will pop up showing 5 digit pictures with their true and predicted labels.
3. Celebrate!:
- If it works, you’ve just trained an AI model! A 95% accuracy means it’s pretty good at reading digits, though it might miss some tricky ones. Thanks for reading, More awesome blogs are on the way with SightSpeak AI, so stay tuned for what’s next!