LeNet-5网络pytorch实现 测试CIFAR10数据集

重写Dataset

没有用官方导入的数据集，而是自己读取图片和标签，主要熟悉一下处理数据集的流程。
dataset.py

from torch.utils.data import Dataset
from PIL import Image  # 读取图片
import os


# 创建一个class，继承Dataset类
class MyDataset(Dataset):
    def __init__(self, path, transform=None):
        self.path = path
        self.transform = transform
        self.image_path = os.path.join(self.path, 'images')  # 获得图片的路径
        self.label_path = os.path.join(self.path, 'labels')  # 获得标签的路径
        self.image_list = os.listdir(self.image_path)  # 获得图片列表
        self.label_list = os.listdir(self.label_path)  # 获得标签列表

    def __getitem__(self, idx):  # idx为编号
        # 获取图片
        img_name = self.image_list[idx]  # 名称
        img_item_path = os.path.join(self.image_path, img_name)
        img = Image.open(img_item_path).convert('RGB')  # 读取图片
        if self.transform is not None:
            img = self.transform(img)
        # 获取标签
        label_name = self.label_list[idx]  # 名称
        label_item_path = os.path.join(self.label_path, label_name)
        with open(label_item_path, 'r') as f:
            label = int(f.read())
        return img, label

    def __len__(self):  # 数据集的长度
        return len(self.image_list)


# # 用类创建实例
# path = 'data/cifar10/train/'
# dataset = MyDataset(path)
# img, label = dataset[0]
# img.show()
# print(label)
# a = 1

Model

model.py

# model.py

import torch.nn as nn
import torch.nn.functional as F


class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 5)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(16, 32, 5)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(32 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))  # input(3,32,32) output(16,28,28)
        x = self.pool1(x)  # output(16，14，14)
        x = F.relu(self.conv2(x))  # output(32,10.10)
        x = self.pool2(x)  # output(32,5,5)
        x = x.view(-1, 32 * 5 * 5)  # output(5*5*32)
        x = F.relu(self.fc1(x))  # output(120)
        x = F.relu(self.fc2(x))  # output(84)
        x = self.fc3(x)  # output(10)
        return x

# #model调试
# import torch
#
# #定义shape
# input1 = torch.rand([36,3,32,32])
# print(input1)
# model = LeNet()#实例化
# print(model)
# #输入网络中
# output = model(input1)
# print(output)
# a = 1

Train

train.py

# train.py

import torch
import torchvision
import torch.nn as nn
from dataset import MyDataset
from model import LeNet
import torch.optim as optim
import torchvision.transforms as transforms
from torchvision import transforms, datasets, utils
import matplotlib.pyplot as plt
import numpy as np


# device : GPU or CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

# 图像变换
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# 训练图片
train_path = 'data/cifar10/train/'
train_set = MyDataset(train_path, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=36,
                                           shuffle=True, num_workers=0)

# 验证图片
val_set = MyDataset(train_path, transform=transform)
val_loader = torch.utils.data.DataLoader(val_set, batch_size=50,
                                         shuffle=True, num_workers=0)
val_data_iter = iter(val_loader)
val_image, val_label = val_data_iter.next()
print(val_image.size())

# 网络
net = LeNet()
net.to(device)
# 损失函数
loss_function = nn.CrossEntropyLoss()
# 定义优化器
optimizer = optim.Adam(net.parameters(), lr=0.001)


# 训练过程
for epoch in range(10):  # loop over the dataset multiple times
    running_loss = 0.0  # 累加损失
    for step, data in enumerate(train_loader, start=0):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data
        # print(inputs.size(), labels.size())
        # zero the parameter gradients
        optimizer.zero_grad() # 如果不清除历史梯度，就会对计算的历史梯度进行累加
        # forward + backward + optimize
        outputs = net(inputs.to(device))
        loss = loss_function(outputs, labels.to(device))
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.item()
        if step % 500 == 499:    # print every 500 mini-batches
            with torch.no_grad():  # 上下文管理器
                outputs = net(val_image.to(device))  # [batch, 10]
                predict_y = torch.max(outputs, dim=1)[1]
                accuracy = (predict_y == val_label.to(device)).sum().item() / val_label.size(0)
                print('[%d, %5d] train_loss: %.3f  test_accuracy: %.3f' %
                      (epoch + 1, step + 1, running_loss / 500, accuracy))
                running_loss = 0.0

print('Finished Training')
# 保存模型
save_path = './Lenet.pth'
torch.save(net.state_dict(), save_path)

训练过程

Test

test.py

# test.py

import torch
import torchvision.transforms as transforms
from PIL import Image

from dataset import MyDataset
from model import LeNet
import matplotlib.pyplot as plt
import numpy as np

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 预处理
transform = transforms.Compose(
    [transforms.Resize((32, 32)),
     transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# 测试图片
test_path = 'data/cifar10/test/'
test_set = MyDataset(test_path, transform=transform)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=36,
                                           shuffle=True, num_workers=0)

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

net = LeNet()
net.load_state_dict(torch.load('./Lenet.pth'))
net.to(device)

correct = 0
total = 0
for step, data in enumerate(test_loader, start=0):
    with torch.no_grad():
        inputs, labels = data
        inputs = inputs.to(device)
        labels = labels.to(device)
        outputs = net(inputs.to(device))
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('Test Accuracy of the model on the 10000 test images: {} %'.format(100 * correct / total))

测试集结果

参考资料

CNN经典网络模型详解-LeNet-5（pytorch实现）