LeNet pytorch实现

数据集MNIST
有一部分代码是直接拷过来的(不擅长数据处理部分)
新建一个DataTest包,里面创建一个MNIST.py

import torch.utils.data
from torchvision import datasets as dataset
from torchvision import transforms


def getTransforms():
    transform = transforms.Compose(
        [transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3018,))]
    )
    return transform


def get_train_dataset():
    return dataset.MNIST(
        root='./MNIST',
        train=True,
        download=True,
        transform=getTransforms()
    )


def get_test_dataset():
    return dataset.MNIST(
        root='./MNIST',
        train=False,
        download=True,
        transform=getTransforms()
    )


def get_train_loader(batch_size, shuffle=True):
    return torch.utils.data.DataLoader(
        dataset=get_train_dataset(),
        batch_size=batch_size,
        shuffle=shuffle
    )


def get_test_loader(batch_size, shuffle=True):
    return torch.utils.data.DataLoader(
        dataset=get_test_dataset(),
        batch_size=batch_size,
        shuffle=shuffle
    )

新建一个LeNet.py(这个名字应该没有要求,随便了)

from torch import nn
import torch
from DataTest import MNIST
import numpy as np


lr = 0.01
momentum = 0.5
epochs = 10
batch_size = 64
test_batch_size = 1000
log_interval = 10  # 跑多少次batch进行一次日志记录


class LeNet(nn.Module):
    def __init__(self):
        # 直接按输入为28*28来计算
        super(LeNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Flatten()
        )
        self.layer3 = nn.Sequential(
            nn.Linear(in_features=16 * 5 * 5, out_features=120),
            nn.ReLU()
        )
        self.layer4 = nn.Sequential(
            nn.Linear(in_features=120, out_features=84),
            nn.ReLU()
        )
        self.layerOutput = nn.Linear(84, 10)

    def forward(self, X):
        x = self.layer1(X)
        x = self.layer2(x)
        # 其他地方见到的展平操作
        # x = x.view(x.size()[0],-1)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.layerOutput(x)
        return x


def train(model, train_loader, epoch, device, trainer):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data = data.to(device)
        target = target.to(device)
        trainer.zero_grad()
        output = model(data)
        loss = nn.CrossEntropyLoss()
        loss = loss(output, target)
        loss.backward()
        trainer.step()
        if batch_idx % log_interval == 0:
            # 打印相关信息
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss:{:.6F}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()
            ))


def test(model, test_loader, device):
    model.eval()
    test_loss = 0
    correct = 0
    for data, target in test_loader:
        data = data.to(device)
        target = target.to(device)
        output = model(data)
        # sum up batch loss 把所有loss值进行累加
        loss = nn.CrossEntropyLoss(reduction='sum')
        test_loss += loss(output, target).item()
        pred = output.data.max(1, keepdim=True)[1]
        # get the index of the max log-probability
        correct += pred.eq(target.data.view_as(pred)).cpu().sum()
        # 对预测正确的数据个数进行累加

    test_loss /= len(test_loader.dataset)
    # 因为把所有loss值进行过累加，所以最后要除以总得数据长度才得平均loss
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))


def run(device):
    model = LeNet()
    model = model.to(device)
    train_loader = MNIST.get_train_loader(batch_size=batch_size)
    test_loader = MNIST.get_test_loader(batch_size=test_batch_size)
    trainer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum)

    for epoch in range(1, epochs + 1):
        train(model, train_loader, epoch, device, trainer)
        test(model, test_loader, device)

    torch.save(model, 'LeNet.pth')


if __name__ == "__main__":
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    run(device)

运行结果:

参考资料:
https://blog.csdn.net/wang_xinyu/article/details/114670155

因为是纯粹的个人练习,所以有些地方可能不严谨(比如激活函数什么的)