LeNet复现-Pytorch版本

模型

from torch import nn


class Lenet(nn.Module):  # 继承Moudle，后面没有括号
    def __init__(self, in_dim, n_class):
        super().__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=in_dim, out_channels=20, kernel_size=5, stride=1, padding=2),  # (20, 28, 28)
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)  # (20, 14, 14) 池化不会改变通道数
        )

        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=20, out_channels=50, kernel_size=5, stride=1, padding=2),  # (50, 14, 14)
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)  # (50, 7, 7)
        )

        # fully_connected 两层
        self.fc = nn.Sequential(
            nn.Linear(in_features=2450, out_features=500),  # (50 * 7 * 7, 500) 降维
            nn.ReLU(),
            nn.Linear(in_features=500, out_features=n_class)  # (500, 10)
        )

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        # flatten之后送入全连接
        x = x.view(x.size(0), -1)
        output = self.fc(x)  # 不需要softmax，因为CrossEntropyLoss做了softmax的处理

        return output

训练

import torch
from torch import nn
from LeNet import Lenet
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


def get_datasets():
    train_dataset = datasets.MNIST(
        root='./data/mnist',
        train=True,
        transform=transforms.ToTensor(),
        download=True,
    )
    test_dataset = datasets.MNIST(
        root='./data/mnist',
        train=False,
        transform=transforms.ToTensor()
    )
    train_loader = DataLoader(
        dataset=train_dataset,
        batch_size=batch_size,
        shuffle=True
    )
    test_loader = DataLoader(
        dataset=test_dataset,
        batch_size=batch_size,
        shuffle=False
    )

    return train_loader, test_loader


def evaluate_accuracy(data_iter, model):
    total = 0
    correct = 0
    with torch.no_grad():
        for images, labels in data_iter:
            model.eval()
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            _, predicts = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicts == labels).cpu().sum()

    return 100 * correct / total


def train(learning_rate, num_epochs, train_loader, test_loader, model):
    # 设置随机种子,确保结果可重复
    torch.manual_seed(1)
    # Adam optimizer
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(num_epochs):
        print('current epoch = %d' % epoch)
        train_accuracy_total = 0
        train_correct = 0
        train_loss_sum = 0
        for i, (images, labels) in enumerate(train_loader):
            train_accuracy_total = 0
            train_correct = 0
            train_loss_sum = 0
            model.train()
            images = images.to(device)
            labels = labels.to(device)
            # forward
            outputs = model(images)
            # calculate current iteration's loss
            loss = criterion(outputs, labels)
            # reset grad graph
            optimizer.zero_grad()
            # backward -
            loss.backward()
            # update parameters
            optimizer.step()
            train_loss_sum += loss.item()
            _, predicts = torch.max(outputs.data, dim=1)  # dim=1 表示按照行求最大值,和其索引
            train_accuracy_total += labels.size(0)
            train_correct += (predicts == labels).cpu().sum().item()
        test_acc = evaluate_accuracy(test_loader, model)
        print('epoch %d, loss %.4f, train accuracy %.3f,'
              ' test accuracy %.3f' %
              (epoch, train_loss_sum / batch_size, train_correct / train_accuracy_total, test_acc))
    torch.save(model, 'lenet.pth')
    print('finish training')


def main():
    model = Lenet(1, 10)  # 单通道灰度图像,数字图像10分类问题,实例化模型
    model.to(device)  # 放到gpu上
    learning_rate = 1e-3  # 学习率
    num_epochs = 10  # epoch num
    # 获取数据集
    train_loader, test_loader = get_datasets()
    # train
    train(learning_rate=learning_rate, num_epochs=num_epochs, train_loader=train_loader, test_loader=test_loader,
          model=model)


if __name__ == '__main__':
    # global variables
    batch_size = 64
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    main()

单张预测：

import torch
from LeNet import Lenet
from PIL import Image
from torchvision import datasets, transforms


def main():
    # 使用gpu训练的话，模型的参数类型是torch.cuda.FloatTensor所以需要to device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 这里还有一个中括号
    transform = transforms.Compose([
        transforms.Resize((28, 28)),
        transforms.ToTensor()
    ])

    test_img = Image.open('number.jpeg')
    test_img = test_img.convert('L')
    print(test_img.size)  # [1115, 1115]
    test_img = transform(test_img)
    test_img = torch.unsqueeze(test_img, dim=0)
    print(test_img.size())  # [1, 1, 28, 28]
    model = torch.load('./lenet.pth')
    model.eval()

    with torch.no_grad():

        test_img = test_img.to(device)  # 使用Gpu训练得到的模型参数是
        print('test_img data type', test_img.dtype)
        output = model(test_img)
        print(output.dtype)
        _, predict = torch.max(output, dim=1)
        print(type(predict))
        print(predict)  # cuda tensor
        print(predict.cpu())  # tensor
        print('预测为数字 ', predict.item())


if __name__ == '__main__':
    main()