LeNet简单实现

1 LeNet

import torch.nn as nn


class LeNet(nn.Module):
    def __init__(self, num_classes=10):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5, stride=1)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5, stride=1)
        self.relu = nn.ReLU(inplace=True)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, num_classes)

    def forward(self, x):
        # input shape: (batch_size, in_channels, 32, 32)
        x = self.conv1(x)
        # shape: (batch_size, 6, 28, 28)
        x = self.relu(x)
        # shape: (batch_size, 6, 28, 28)
        x = self.pool(x)
        # shape: (batch_size, 6, 14, 14)
        x = self.conv2(x)
        # shape: (batch_size, 16, 10, 10)
        x = self.relu(x)
        # shape: (batch_size, 16, 10, 10)
        x = self.pool(x)
        # shape: (batch_size, 16, 5, 5)
        x = x.view(x.size(0), -1)
        # shape: (batch_size, 16 * 5 * 5)
        x = self.fc1(x)
        # shape: (batch_size, 120)
        x = self.relu(x)
        # shape: (batch_size, 120)
        x = self.fc2(x)
        # shape: (batch_size, 84)
        x = self.relu(x)
        # shape: (batch_size, 84)
        x = self.fc3(x)
        # shape: (batch_size, num_classes)
        return x

2 Train

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision.datasets import MNIST
import torchvision.transforms as transforms
from torchvision.transforms import ToTensor
from torch.utils.data import DataLoader
from model import LeNet

# Define the transform to normalize the data
transform = transforms.Compose([
    transforms.Resize((32,32)),
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# Define hyperparameters
batch_size = 64
learning_rate = 0.01
num_epochs = 10

# Load the MNIST dataset
train_dataset = MNIST(root='./dataset', train=True, download=True, transform=transform)
test_dataset = MNIST(root='./dataset', train=False, download=True, transform=transform)

# Create data loaders
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# Create the LeNet model
model = LeNet()

# Define the loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)

# Train the model
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)

        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Print training loss and accuracy
        if (i+1) % 100 == 0:
            total = labels.size(0)
            _, predicted = torch.max(outputs.data, 1)
            correct = (predicted == labels).sum().item()
            accuracy = correct / total
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}, Accuracy: {:.2f}%'
                  .format(epoch+1, num_epochs, i+1, len(train_loader), loss.item(), accuracy*100))

# Test the model
model.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    print('Test Accuracy: {:.2f}%'.format(correct / total * 100))