pytorch实现LeNet模型MNIST手写识别

模块导入，常量参数设定

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torchsummary import summary

BATCH_SIZE = 512
EPOCH = 20
DEIVCE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

数据加载

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        'data', 
        train=True, 
        download=False,
        transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])),
    batch_size=BATCH_SIZE, shuffle=True)


test_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        'data', 
        train=False, 
        transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])),
    batch_size=BATCH_SIZE, shuffle=True)

LeNet模型搭建

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        # nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0,dilation=1,groups=1, bias=True）
        self.conv1 = nn.Conv2d(1, 6, 5)  
        self.pool1 = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(6, 16, 3)
        self.pool2 = nn.MaxPool2d(2)
        # nn.Linear(in_features, out_features, bias=True)
        self.fc1   = nn.Linear(16*5*5, 120)
        self.fc2   = nn.Linear(120, 84)
        self.fc3   = nn.Linear(84, 10)
        
    def forward(self, x):
        out = F.relu(self.conv1(x))  # 24
        # nn.MaxPool2d(kernel_size, stride=None, padding=0, dilation=1,return_indices=False,ceil_mode=False)
        # max_pool2d(*args, **kwargs)
        out = self.pool1(out)  # 12
        out = F.relu(self.conv2(out))  # 10
        out = self.pool2(out)  # 5
        out = out.view(out.size(0), -1)  # flatten
        out = F.relu(self.fc1(out))
        out = F.relu(self.fc2(out))
        out = self.fc3(out)
        out = F.log_softmax(out,dim=1)
        return out

查看模型结构

model = LeNet()
summary(model, (1,28,28))  

模型训练测试函数

model = LeNet().to(device)
optimizer = optim.Adam(model.parameters())

def train(device, model, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()  # zero all gradients
        output = model(data)
        loss = F.nll_loss(output, target)  # calculate loss
        loss.backward()
        optimizer.step()  # update all parameters
        if (batch_idx + 1) % 30 == 0:
            print(f"Train Epoch {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} Loss: {loss.item():.6f}]")


def test(device, model, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)  # predict class probability
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # superposition whole batch loss
            pred = output.max(1, keepdim=True)[1]  # return max probability index
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    print(f'\nTest Epoch: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({100. * correct / len(test_loader.dataset):.2f}%)\n')

模型训练，测试

for epoch in range(1, EPOCH + 1):
    train(DEIVCE, model, train_loader, optimizer, epoch)
    test(DEIVCE, model, train_loader)