一、观察MNIST数据集

import numpy as np
import matplotlib.pyplot as plt
import PIL.Image as Image
import json
import gzip

# 设置中文字体为黑体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
# --------------------------------------------------打印并观察数据分布------------------------------------------------------

# 打印并观察数据集分布情况
train_set, dev_set, test_set = json.load(gzip.open('./mnist.json.gz'))
train_images, train_labels = train_set[0][:1000], train_set[1][:1000]
dev_images, dev_labels = dev_set[0][:200], dev_set[1][:200]
test_images, test_labels = test_set[0][:200], test_set[1][:200]
train_set, dev_set, test_set = [train_images, train_labels], [dev_images, dev_labels], [test_images, test_labels]
print('Length of train/dev/test set:{}/{}/{}'.format(len(train_set[0]), len(dev_set[0]), len(test_set[0])))

# ---------------------------------------------------观察一张图像---------------------------------------------------------
image, label = train_set[0][0], train_set[1][0]
image, label = np.array(image).astype('float32'), int(label)
# 原始图像数据为长度784的行向量，需要调整为[28,28]大小的图像
image = np.reshape(image, [28, 28])
image *= 255
image = Image.fromarray(image.astype('uint8'), mode='L')
print("The number in the picture is {}".format(label))
plt.figure(figsize=(5, 5))
plt.imshow(image, cmap='gray')
plt.show()

二、自定义算子实现LeNet

import numpy as np
import matplotlib.pyplot as plt
import PIL.Image as Image
import json
import gzip
import torchvision.transforms as transforms
from torch.utils.data import Dataset
import torch
import torch.nn.functional as F
import torch.nn as nn
from torchsummary import summary
# --------------------------------------------------打印并观察数据分布------------------------------------------------------
train_set, dev_set, test_set = json.load(gzip.open('./mnist.json.gz'))  # 读取数据集
# 获取对应图像与标签
train_images, train_labels = train_set[0][:1000], train_set[1][:1000]
dev_images, dev_labels = dev_set[0][:200], dev_set[1][:200]
test_images, test_labels = test_set[0][:200], test_set[1][:200]
# 划分数据集
train_set, dev_set, test_set = [train_images, train_labels], [dev_images, dev_labels], [test_images, test_labels]
print('Length of train/dev/test set:{}/{}/{}'.format(len(train_set[0]), len(dev_set[0]), len(test_set[0])))
# ---------------------------------------------------观察一张图像---------------------------------------------------------
image, label = train_set[0][0], train_set[1][0]
image, label = np.array(image).astype('float32'), int(label)
# 原始图像数据为长度784的行向量，需要调整为[28,28]大小的图像
image = np.reshape(image, [28, 28])
image *= 255
image = Image.fromarray(image.astype('uint8'), mode='L')
print("The number in the picture is {}".format(label))
plt.figure(figsize=(5, 5))
plt.imshow(image, cmap='gray')
plt.show()
# ---------------------------------------------------数据处理与加载--------------------------------------------------------
# 数据预处理
transforms = transforms.Compose(
    [transforms.Resize(32), transforms.ToTensor(), transforms.Normalize(mean=[0.5], std=[0.5])])


class MNIST_dataset(Dataset):
    def __init__(self, dataset, transforms, mode='train'):
        self.mode = mode
        self.transforms = transforms
        self.dataset = dataset

    def __getitem__(self, idx):
        # 获取图像和标签
        image, label = self.dataset[0][idx], self.dataset[1][idx]
        image, label = np.array(image).astype('float32'), int(label)
        image = np.reshape(image, [28, 28])
        image = Image.fromarray(image.astype('uint8'), mode='L')
        image = self.transforms(image)
        return image, label

    def __len__(self):
        return len(self.dataset[0])


# 加载 mnist 数据集
train_dataset = MNIST_dataset(dataset=train_set, transforms=transforms, mode='train')
test_dataset = MNIST_dataset(dataset=test_set, transforms=transforms, mode='test')
dev_dataset = MNIST_dataset(dataset=dev_set, transforms=transforms, mode='dev')


# -----------------------------------------------------定义算子 ---------------------------------------------------------
class Conv2D(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
        super(Conv2D, self).__init__()
        # 初始化卷积核权重和偏置
        self.weight = nn.Parameter(torch.ones(out_channels, in_channels, kernel_size, kernel_size))
        self.bias = nn.Parameter(torch.zeros(out_channels, 1))
        self.stride = stride
        self.padding = padding
        self.in_channels = in_channels
        self.out_channels = out_channels

    # 基础卷积运算
    def single_forward(self, X, weight):
        # 零填充
        new_X = torch.zeros([X.shape[0], X.shape[1] + 2 * self.padding, X.shape[2] + 2 * self.padding])
        new_X[:, self.padding:X.shape[1] + self.padding, self.padding:X.shape[2] + self.padding] = X
        u, v = weight.shape
        output_w = (new_X.shape[1] - u) // self.stride + 1
        output_h = (new_X.shape[2] - v) // self.stride + 1
        output = torch.zeros([X.shape[0], output_w, output_h])
        for i in range(output.shape[1]):
            for j in range(output.shape[2]):
                output[:, i, j] = torch.sum(
                    new_X[:, i * self.stride:i * self.stride + u, j * self.stride:j * self.stride + v] * weight,
                    dim=[1, 2])
        return output

    def forward(self, inputs):
        """
        输入：
            - inputs: 输入矩阵，形状为 [B, D, M, N]
            - weights: P组二维卷积核，形状为 [P, D, U, V]
            - bias: P个偏置，形状为 [P, 1]
        """
        feature_maps = []
        # 进行多次多输入通道卷积运算
        for p in range(self.out_channels):
            multi_outs = []
            for i in range(self.in_channels):
                single = self.single_forward(inputs[:, i, :, :], self.weight[p, i, :, :])
                multi_outs.append(single)
            feature_map = torch.sum(torch.stack(multi_outs), dim=0) + self.bias[p]
            feature_maps.append(feature_map)
        # 将所有特征图堆叠起来
        out = torch.stack(feature_maps, dim=1)
        return out


class Pool2D(nn.Module):
    def __init__(self, size=(2, 2), mode='max', stride=1):
        super(Pool2D, self).__init__()
        self.mode = mode
        self.h, self.w = size
        self.stride = stride

    def forward(self, x):
        # 计算输出大小
        output_w = (x.shape[2] - self.w) // self.stride + 1
        output_h = (x.shape[3] - self.h) // self.stride + 1
        output = torch.zeros([x.shape[0], x.shape[1], output_w, output_h])

        # 进行池化操作
        for i in range(output.shape[2]):
            for j in range(output.shape[3]):
                # 最大池化
                if self.mode == 'max':
                    window = x[:, :, self.stride * i:self.stride * i + self.h, self.stride * j:self.stride * j + self.w]
                    # 最大池化沿着高度和宽度
                    max_values, _ = torch.max(window, dim=2)  # 沿高度方向
                    max_values, _ = torch.max(max_values, dim=2)  # 沿宽度方向
                    output[:, :, i, j] = max_values

                # 平均池化
                elif self.mode == 'avg':
                    window = x[:, :, self.stride * i:self.stride * i + self.h, self.stride * j:self.stride * j + self.w]
                    # 平均池化沿着高度和宽度
                    avg_values = torch.mean(window, dim=[2, 3])
                    output[:, :, i, j] = avg_values

        return output


# -----------------------------------------------------构建网络----------------------------------------------------------
class Model_LeNet(nn.Module):
    def __init__(self, in_channels, num_classes=10):
        super(Model_LeNet, self).__init__()
        # 卷积层：输出通道数为6，卷积核大小为5×5
        self.conv1 = Conv2D(in_channels=in_channels, out_channels=6, kernel_size=5)
        # 汇聚层：汇聚窗口为2×2，步长为2
        self.pool2 = Pool2D(size=(2, 2), mode='max', stride=2)
        # 卷积层：输入通道数为6，输出通道数为16，卷积核大小为5×5，步长为1
        self.conv3 = Conv2D(in_channels=6, out_channels=16, kernel_size=5, stride=1)
        # 汇聚层：汇聚窗口为2×2，步长为2
        self.pool4 = Pool2D(size=(2, 2), mode='avg', stride=2)
        # 卷积层：输入通道数为16，输出通道数为120，卷积核大小为5×5
        self.conv5 = Conv2D(in_channels=16, out_channels=120, kernel_size=5, stride=1)
        # 全连接层：输入神经元为120，输出神经元为84
        self.linear6 = nn.Linear(120, 84)
        # 全连接层：输入神经元为84，输出神经元为类别数
        self.linear7 = nn.Linear(84, num_classes)

    def forward(self, x):
        # C1：卷积层+激活函数
        output = F.relu(self.conv1(x))
        # S2：汇聚层
        output = self.pool2(output)
        # C3：卷积层+激活函数
        output = F.relu(self.conv3(output))
        # S4：汇聚层
        output = self.pool4(output)
        # C5：卷积层+激活函数
        output = F.relu(self.conv5(output))
        # 输入层将数据拉平[B,C,H,W] -> [B,CxHxW]
        output = torch.squeeze(output, dim=3)
        output = torch.squeeze(output, dim=2)
        # F6：全连接层
        output = F.relu(self.linear6(output))
        # F7：全连接层
        output = self.linear7(output)
        return output


# -----------------------------------------------------网络测试-----------------------------------------------------------
# 这里用np.random创建一个随机数组作为输入数据
inputs = np.random.randn(*[1, 1, 32, 32])
inputs = inputs.astype('float32')
model = Model_LeNet(in_channels=1, num_classes=10)
c = []
for a, b in model.named_children():
    c.append(a)
print(c)
x = torch.tensor(inputs)
for a, item in model.named_children():
    try:
        x = item(x)
    except:
        x = torch.reshape(x, [x.shape[0], -1])
        x = item(x)
    d = []
    e = []
    for b, c in item.named_parameters():
        d.append(b)
        e.append(c)
    if len(e) == 2:
        print(a, x.shape, e[0].shape,
              e[1].shape)
    else:
        # 汇聚层没有参数
        print(a, x.shape)

Torch_model = Model_LeNet(in_channels=1, num_classes=10)
summary(Torch_model, (1, 32, 32))

三、预定义算子实现LeNet

import numpy as np
import matplotlib.pyplot as plt
import PIL.Image as Image
import json
import gzip
import torchvision.transforms as transforms
from torch.utils.data import Dataset
import torch
import torch.nn.functional as F
import torch.nn as nn
from torchsummary import summary

# --------------------------------------------------打印并观察数据分布------------------------------------------------------
train_set, dev_set, test_set = json.load(gzip.open('./mnist.json.gz'))  # 读取数据集
# 获取对应图像与标签
train_images, train_labels = train_set[0][:1000], train_set[1][:1000]
dev_images, dev_labels = dev_set[0][:200], dev_set[1][:200]
test_images, test_labels = test_set[0][:200], test_set[1][:200]
# 划分数据集
train_set, dev_set, test_set = [train_images, train_labels], [dev_images, dev_labels], [test_images, test_labels]
print('Length of train/dev/test set:{}/{}/{}'.format(len(train_set[0]), len(dev_set[0]), len(test_set[0])))
# ---------------------------------------------------观察一张图像---------------------------------------------------------
image, label = train_set[0][0], train_set[1][0]
image, label = np.array(image).astype('float32'), int(label)
# 原始图像数据为长度784的行向量，需要调整为[28,28]大小的图像
image = np.reshape(image, [28, 28])
image *= 255
image = Image.fromarray(image.astype('uint8'), mode='L')
print("The number in the picture is {}".format(label))
plt.figure(figsize=(5, 5))
plt.imshow(image, cmap='gray')
plt.show()
# ---------------------------------------------------数据处理与加载--------------------------------------------------------
# 数据预处理
transforms = transforms.Compose(
    [transforms.Resize(32), transforms.ToTensor(), transforms.Normalize(mean=[0.5], std=[0.5])])


class MNIST_dataset(Dataset):
    def __init__(self, dataset, transforms, mode='train'):
        self.mode = mode
        self.transforms = transforms
        self.dataset = dataset

    def __getitem__(self, idx):
        # 获取图像和标签
        image, label = self.dataset[0][idx], self.dataset[1][idx]
        image, label = np.array(image).astype('float32'), int(label)
        image = np.reshape(image, [28, 28])
        image = Image.fromarray(image.astype('uint8'), mode='L')
        image = self.transforms(image)
        return image, label

    def __len__(self):
        return len(self.dataset[0])


# 加载 mnist 数据集
train_dataset = MNIST_dataset(dataset=train_set, transforms=transforms, mode='train')
test_dataset = MNIST_dataset(dataset=test_set, transforms=transforms, mode='test')
dev_dataset = MNIST_dataset(dataset=dev_set, transforms=transforms, mode='dev')


# -----------------------------------------------------构建网络----------------------------------------------------------
class PyTorch_LeNet(nn.Module):
    def __init__(self, in_channels, num_classes=10):
        super(PyTorch_LeNet, self).__init__()
        # 卷积层：输出通道数为6，卷积核大小为5*5
        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=6, kernel_size=5)
        # 汇聚层：汇聚窗口为2*2，步长为2
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        # 卷积层：输入通道数为6，输出通道数为16，卷积核大小为5*5
        self.conv3 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5)
        # 汇聚层：汇聚窗口为2*2，步长为2
        self.pool4 = nn.AvgPool2d(kernel_size=2, stride=2)
        # 卷积层：输入通道数为16，输出通道数为120，卷积核大小为5*5
        self.conv5 = nn.Conv2d(in_channels=16, out_channels=120, kernel_size=5)
        # 全连接层：输入神经元为120，输出神经元为84
        self.linear6 = nn.Linear(in_features=120, out_features=84)
        # 全连接层：输入神经元为84，输出神经元为类别数
        self.linear7 = nn.Linear(in_features=84, out_features=num_classes)

    def forward(self, x):
        # C1：卷积层+激活函数
        output = F.relu(self.conv1(x))
        # S2：汇聚层
        output = self.pool2(output)
        # C3：卷积层+激活函数
        output = F.relu(self.conv3(output))
        # S4：汇聚层
        output = self.pool4(output)
        # C5：卷积层+激活函数
        output = F.relu(self.conv5(output))
        # 输入层将数据拉平[B,C,H,W] -> [B,CxHxW]
        output = torch.squeeze(output, dim=3)
        output = torch.squeeze(output, dim=2)
        # F6：全连接层
        output = F.relu(self.linear6(output))
        # F7：全连接层
        output = self.linear7(output)
        return output


# -----------------------------------------------------网络测试----------------------------------------------------------
# 这里用np.random创建一个随机数组作为输入数据
inputs = np.random.randn(*[1, 1, 32, 32])
inputs = inputs.astype('float32')
model = PyTorch_LeNet(in_channels=1, num_classes=10)
c = []
for a, b in model.named_children():
    c.append(a)
print(c)
x = torch.tensor(inputs)
for a, item in model.named_children():
    try:
        x = item(x)
    except:
        x = torch.reshape(x, [x.shape[0], -1])
        x = item(x)
    d = []
    e = []
    for b, c in item.named_parameters():
        d.append(b)
        e.append(c)
    if len(e) == 2:
        print(a, x.shape, e[0].shape,
              e[1].shape)
    else:
        # 汇聚层没有参数
        print(a, x.shape)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # PyTorch v0.4.0
Torch_model = PyTorch_LeNet(in_channels=1, num_classes=10).to(device)
summary(Torch_model, (1, 32, 32))

四、对比运行时间

import numpy as np
import matplotlib.pyplot as plt
import PIL.Image as Image
import json
import gzip
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
import torch
import torch.nn.functional as F
import torch.nn as nn
import time

# --------------------------------------------------打印并观察数据分布------------------------------------------------------
train_set, dev_set, test_set = json.load(gzip.open('./mnist.json.gz'))  # 读取数据集
# 获取对应图像与标签
train_images, train_labels = train_set[0][:1000], train_set[1][:1000]
dev_images, dev_labels = dev_set[0][:200], dev_set[1][:200]
test_images, test_labels = test_set[0][:200], test_set[1][:200]
# 划分数据集
train_set, dev_set, test_set = [train_images, train_labels], [dev_images, dev_labels], [test_images, test_labels]
print('Length of train/dev/test set:{}/{}/{}'.format(len(train_set[0]), len(dev_set[0]), len(test_set[0])))
# ---------------------------------------------------观察一张图像---------------------------------------------------------
image, label = train_set[0][0], train_set[1][0]
image, label = np.array(image).astype('float32'), int(label)
# 原始图像数据为长度784的行向量，需要调整为[28,28]大小的图像
image = np.reshape(image, [28, 28])
image *= 255
image = Image.fromarray(image.astype('uint8'), mode='L')
print("The number in the picture is {}".format(label))
plt.figure(figsize=(5, 5))
plt.imshow(image, cmap='gray')
plt.show()
# ---------------------------------------------------数据处理与加载--------------------------------------------------------
# 数据预处理
transforms = transforms.Compose(
    [transforms.Resize(32), transforms.ToTensor(), transforms.Normalize(mean=[0.5], std=[0.5])])


class MNIST_dataset(Dataset):
    def __init__(self, dataset, transforms, mode='train'):
        self.mode = mode
        self.transforms = transforms
        self.dataset = dataset

    def __getitem__(self, idx):
        # 获取图像和标签
        image, label = self.dataset[0][idx], self.dataset[1][idx]
        image, label = np.array(image).astype('float32'), int(label)
        image = np.reshape(image, [28, 28])
        image = Image.fromarray(image.astype('uint8'), mode='L')
        image = self.transforms(image)
        return image, label

    def __len__(self):
        return len(self.dataset[0])


# 加载 mnist 数据集
train_dataset = MNIST_dataset(dataset=train_set, transforms=transforms, mode='train')
test_dataset = MNIST_dataset(dataset=test_set, transforms=transforms, mode='test')
dev_dataset = MNIST_dataset(dataset=dev_set, transforms=transforms, mode='dev')


# -----------------------------------------------------定义算子 ---------------------------------------------------------
class Conv2D(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
        super(Conv2D, self).__init__()
        # 初始化卷积核权重和偏置
        self.weight = nn.Parameter(torch.ones(out_channels, in_channels, kernel_size, kernel_size))
        self.bias = nn.Parameter(torch.zeros(out_channels, 1))
        self.stride = stride
        self.padding = padding
        self.in_channels = in_channels
        self.out_channels = out_channels

    # 基础卷积运算
    def single_forward(self, X, weight):
        # 零填充
        new_X = torch.zeros([X.shape[0], X.shape[1] + 2 * self.padding, X.shape[2] + 2 * self.padding])
        new_X[:, self.padding:X.shape[1] + self.padding, self.padding:X.shape[2] + self.padding] = X
        u, v = weight.shape
        output_w = (new_X.shape[1] - u) // self.stride + 1
        output_h = (new_X.shape[2] - v) // self.stride + 1
        output = torch.zeros([X.shape[0], output_w, output_h])
        for i in range(output.shape[1]):
            for j in range(output.shape[2]):
                output[:, i, j] = torch.sum(
                    new_X[:, i * self.stride:i * self.stride + u, j * self.stride:j * self.stride + v] * weight,
                    dim=[1, 2])
        return output

    def forward(self, inputs):
        """
        输入：
            - inputs: 输入矩阵，形状为 [B, D, M, N]
            - weights: P组二维卷积核，形状为 [P, D, U, V]
            - bias: P个偏置，形状为 [P, 1]
        """
        feature_maps = []
        # 进行多次多输入通道卷积运算
        for p in range(self.out_channels):
            multi_outs = []
            for i in range(self.in_channels):
                single = self.single_forward(inputs[:, i, :, :], self.weight[p, i, :, :])
                multi_outs.append(single)
            feature_map = torch.sum(torch.stack(multi_outs), dim=0) + self.bias[p]
            feature_maps.append(feature_map)
        # 将所有特征图堆叠起来
        out = torch.stack(feature_maps, dim=1)
        return out


class Pool2D(nn.Module):
    def __init__(self, size=(2, 2), mode='max', stride=1):
        super(Pool2D, self).__init__()
        self.mode = mode
        self.h, self.w = size
        self.stride = stride

    def forward(self, x):
        # 计算输出大小
        output_w = (x.shape[2] - self.w) // self.stride + 1
        output_h = (x.shape[3] - self.h) // self.stride + 1
        output = torch.zeros([x.shape[0], x.shape[1], output_w, output_h])

        # 进行池化操作
        for i in range(output.shape[2]):
            for j in range(output.shape[3]):
                # 最大池化
                if self.mode == 'max':
                    window = x[:, :, self.stride * i:self.stride * i + self.h, self.stride * j:self.stride * j + self.w]
                    # 最大池化沿着高度和宽度
                    max_values, _ = torch.max(window, dim=2)  # 沿高度方向
                    max_values, _ = torch.max(max_values, dim=2)  # 沿宽度方向
                    output[:, :, i, j] = max_values

                # 平均池化
                elif self.mode == 'avg':
                    window = x[:, :, self.stride * i:self.stride * i + self.h, self.stride * j:self.stride * j + self.w]
                    # 平均池化沿着高度和宽度
                    avg_values = torch.mean(window, dim=[2, 3])
                    output[:, :, i, j] = avg_values

        return output


# -----------------------------------------------------构建网络----------------------------------------------------------
class Model_LeNet(nn.Module):
    def __init__(self, in_channels, num_classes=10):
        super(Model_LeNet, self).__init__()
        # 卷积层：输出通道数为6，卷积核大小为5×5
        self.conv1 = Conv2D(in_channels=in_channels, out_channels=6, kernel_size=5)
        # 汇聚层：汇聚窗口为2×2，步长为2
        self.pool2 = Pool2D(size=(2, 2), mode='max', stride=2)
        # 卷积层：输入通道数为6，输出通道数为16，卷积核大小为5×5，步长为1
        self.conv3 = Conv2D(in_channels=6, out_channels=16, kernel_size=5, stride=1)
        # 汇聚层：汇聚窗口为2×2，步长为2
        self.pool4 = Pool2D(size=(2, 2), mode='avg', stride=2)
        # 卷积层：输入通道数为16，输出通道数为120，卷积核大小为5×5
        self.conv5 = Conv2D(in_channels=16, out_channels=120, kernel_size=5, stride=1)
        # 全连接层：输入神经元为120，输出神经元为84
        self.linear6 = nn.Linear(120, 84)
        # 全连接层：输入神经元为84，输出神经元为类别数
        self.linear7 = nn.Linear(84, num_classes)

    def forward(self, x):
        # C1：卷积层+激活函数
        output = F.relu(self.conv1(x))
        # S2：汇聚层
        output = self.pool2(output)
        # C3：卷积层+激活函数
        output = F.relu(self.conv3(output))
        # S4：汇聚层
        output = self.pool4(output)
        # C5：卷积层+激活函数
        output = F.relu(self.conv5(output))
        # 输入层将数据拉平[B,C,H,W] -> [B,CxHxW]
        output = torch.squeeze(output, dim=3)
        output = torch.squeeze(output, dim=2)
        # F6：全连接层
        output = F.relu(self.linear6(output))
        # F7：全连接层
        output = self.linear7(output)
        return output


class PyTorch_LeNet(nn.Module):
    def __init__(self, in_channels, num_classes=10):
        super(PyTorch_LeNet, self).__init__()
        # 卷积层：输出通道数为6，卷积核大小为5*5
        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=6, kernel_size=5)
        # 汇聚层：汇聚窗口为2*2，步长为2
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        # 卷积层：输入通道数为6，输出通道数为16，卷积核大小为5*5
        self.conv3 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5)
        # 汇聚层：汇聚窗口为2*2，步长为2
        self.pool4 = nn.AvgPool2d(kernel_size=2, stride=2)
        # 卷积层：输入通道数为16，输出通道数为120，卷积核大小为5*5
        self.conv5 = nn.Conv2d(in_channels=16, out_channels=120, kernel_size=5)
        # 全连接层：输入神经元为120，输出神经元为84
        self.linear6 = nn.Linear(in_features=120, out_features=84)
        # 全连接层：输入神经元为84，输出神经元为类别数
        self.linear7 = nn.Linear(in_features=84, out_features=num_classes)

    def forward(self, x):
        # C1：卷积层+激活函数
        output = F.relu(self.conv1(x))
        # S2：汇聚层
        output = self.pool2(output)
        # C3：卷积层+激活函数
        output = F.relu(self.conv3(output))
        # S4：汇聚层
        output = self.pool4(output)
        # C5：卷积层+激活函数
        output = F.relu(self.conv5(output))
        # 输入层将数据拉平[B,C,H,W] -> [B,CxHxW]
        output = torch.squeeze(output, dim=3)
        output = torch.squeeze(output, dim=2)
        # F6：全连接层
        output = F.relu(self.linear6(output))
        # F7：全连接层
        output = self.linear7(output)
        return output


# -----------------------------------------------------速度比对----------------------------------------------------------
# 这里用np.random创建一个随机数组作为测试数据
inputs = np.random.randn(*[1, 1, 32, 32])
inputs = inputs.astype('float32')
x = torch.tensor(inputs)

# 创建Model_LeNet类的实例，指定模型名称和分类的类别数目
model = Model_LeNet(in_channels=1, num_classes=10)
# 创建Paddle_LeNet类的实例，指定模型名称和分类的类别数目
torch_LeNet = PyTorch_LeNet(in_channels=1, num_classes=10)

# 计算Model_LeNet类的运算速度
model_time = 0
for i in range(60):
    strat_time = time.time()
    out = model(x)
    end_time = time.time()
    # 预热10次运算，不计入最终速度统计
    if i < 10:
        continue
    model_time += (end_time - strat_time)
avg_model_time = model_time / 50
print('Model_LeNet speed:', avg_model_time, 's')

# 计算Paddle_LeNet类的运算速度
torch_model_time = 0
for i in range(60):
    strat_time = time.time()
    torch_out = torch_LeNet(x)
    end_time = time.time()
    # 预热10次运算，不计入最终速度统计
    if i < 10:
        continue
    torch_model_time += (end_time - strat_time)
avg_torch_model_time = torch_model_time / 50
print('PyTorch_LeNet speed:', avg_torch_model_time, 's')
# -----------------------------------------------------同样权重----------------------------------------------------------
# 这里用np.random创建一个随机数组作为测试数据
inputs = np.random.randn(*[1, 1, 32, 32])
inputs = inputs.astype('float32')
x = torch.tensor(inputs)
# 创建Model_LeNet类的实例，指定模型名称和分类的类别数目
model = Model_LeNet(in_channels=1, num_classes=10)
# 获取网络的权重
params = model.state_dict()
# 自定义Conv2D算子的bias参数形状为[out_channels, 1]
# 需要进行调整后才可以赋值
for key in params:
    if 'bias' in key:
        params[key] = params[key].squeeze()
# 创建Paddle_LeNet类的实例，指定模型名称和分类的类别数目
Torch_model = PyTorch_LeNet(in_channels=1, num_classes=10)
# 将Model_LeNet的权重参数赋予给Paddle_LeNet模型，保持两者一致
Torch_model.load_state_dict(params)

# 打印结果保留小数点后6位
torch.set_printoptions(6)
# 计算Model_LeNet的结果
output = model(x)
print('Model_LeNet output: ', output)
# 计算Paddle_LeNet的结果
paddle_output = Torch_model(x)
print('Paddle_LeNet output: ', paddle_output)

五、使用LeNet实现手写数字识别

训练程序

import matplotlib.pyplot as plt
from torchvision import datasets
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, random_split
import torch
import torch.nn.functional as F
import torch.nn as nn
import torch.optim as opt
from RunnerV3 import RunnerV3
from Accuracy import Accuracy

# 设置中文字体为黑体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
# --------------------------------------------------数据处理------------------------------------------------------
# 定义数据预处理：将图像大小调整为32x32，并将其转换为Tensor
transform = transforms.Compose([
    transforms.Resize(32),
    transforms.ToTensor(),
])

# --------------------------------------------------加载MNIST数据集--------------------------------------------------
# 下载并加载 MNIST 数据集
train_set = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_set = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

# --------------------------------------------------划分验证集-----------------------------------------------------
# 划分训练集为训练集和验证集（80%训练集, 20%验证集）
train_size = int(0.8 * len(train_set))  # 80% 用于训练
val_size = len(train_set) - train_size  # 剩余的 20% 用于验证
train_subset, val_subset = random_split(train_set, [train_size, val_size])

# --------------------------------------------------数据加载器-----------------------------------------------------
# 创建 DataLoader
train_loader = DataLoader(train_subset, batch_size=64, shuffle=True)
dev_loader = DataLoader(val_subset, batch_size=64, shuffle=False)
test_loader = DataLoader(test_set, batch_size=64, shuffle=False)

# ---------------------------------------------------检查数据加载------------------------------------------------------
# 检查训练集、验证集和测试集的大小
print(f"Train set size: {len(train_loader.dataset)}")
print(f"Validation set size: {len(dev_loader.dataset)}")
print(f"Test set size: {len(test_loader.dataset)}")

# ---------------------------------------------------数据处理与加载--------------------------------------------------------
# 数据预处理
transforms = transforms.Compose(
    [transforms.Resize(32), transforms.ToTensor()])  # 数据集已经标准化了


# -----------------------------------------------------构建网络----------------------------------------------------------
class PyTorch_LeNet(nn.Module):
    def __init__(self, in_channels, num_classes=10):
        super(PyTorch_LeNet, self).__init__()
        # 卷积层：输出通道数为6，卷积核大小为5*5
        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=6, kernel_size=5)
        # 汇聚层：汇聚窗口为2*2，步长为2
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        # 卷积层：输入通道数为6，输出通道数为16，卷积核大小为5*5
        self.conv3 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5)
        # 汇聚层：汇聚窗口为2*2，步长为2
        self.pool4 = nn.AvgPool2d(kernel_size=2, stride=2)
        # 卷积层：输入通道数为16，输出通道数为120，卷积核大小为5*5
        self.conv5 = nn.Conv2d(in_channels=16, out_channels=120, kernel_size=5)
        # 全连接层：输入神经元为120，输出神经元为84
        self.linear6 = nn.Linear(in_features=120, out_features=84)
        # 全连接层：输入神经元为84，输出神经元为类别数
        self.linear7 = nn.Linear(in_features=84, out_features=num_classes)

    def forward(self, x):
        # C1：卷积层+激活函数
        output = F.relu(self.conv1(x))
        # S2：汇聚层
        output = self.pool2(output)
        # C3：卷积层+激活函数
        output = F.relu(self.conv3(output))
        # S4：汇聚层
        output = self.pool4(output)
        # C5：卷积层+激活函数
        output = F.relu(self.conv5(output))
        # 输入层将数据拉平[B,C,H,W] -> [B,CxHxW]
        output = torch.squeeze(output, dim=3)
        output = torch.squeeze(output, dim=2)
        # F6：全连接层
        output = F.relu(self.linear6(output))
        # F7：全连接层
        output = self.linear7(output)
        return output


# -----------------------------------------------------模型训练----------------------------------------------------------
epoch_num = 10
lr = 0.1
log_steps = 15
eval_steps = 15
batch_size = 64
model_saved_dir = 'C:/Users/chong/Desktop/20221205013/Python/Py/Deep Learning/Train_logs/best_model.pth'
model = PyTorch_LeNet(in_channels=1, num_classes=10)
# 定义优化器
optimizer = opt.SGD(lr=lr, params=model.parameters())
# 定义损失函数
loss_fn = F.cross_entropy
# 定义评价指标
metric = Accuracy(is_logist=True)
# 实例化 RunnerV3 类，并传入训练配置。
runner = RunnerV3(model, optimizer, metric, loss_fn)
runner.train(train_loader, dev_loader, num_epochs=epoch_num, log_steps=log_steps,
             eval_steps=eval_steps,
             save_dir=model_saved_dir)
# ------------------------------------------------------模型测试----------------------------------------------------------
runner.load_model(model_saved_dir)  # 加载训练好的模型
score, loss = runner.evaluate(test_loader)  # 在测试集上对模型进行评价
print("[Test] score/loss: {:.4f}/{:.4f}".format(score, loss))


# ------------------------------------------------------训练过程----------------------------------------------------------

# 绘制训练集和验证集的损失变化以及验证集上的准确率变化曲线
def plot_training_loss_acc(runner, fig_name,
                           fig_size=(16, 6),
                           sample_step=20,
                           loss_legend_loc="upper right",
                           acc_legend_loc="lower right",
                           train_color="#4169E1",
                           val_color='#00BFFF',
                           fontsize='large',
                           train_linestyle="-",
                           val_linestyle='--'):
    plt.figure(figsize=fig_size)
    plt.subplot(1, 2, 1)
    train_items = runner.train_step_losses[::sample_step]
    train_steps = [x[0] for x in train_items]
    train_losses = [x[1] for x in train_items]
    plt.plot(train_steps, train_losses, color=train_color, linestyle=train_linestyle, label="训练损失")
    if len(runner.val_losses) > 0:
        val_steps = [x[0] for x in runner.val_losses]
        val_losses = [x[1] for x in runner.val_losses]
        plt.plot(val_steps, val_losses, color=val_color, linestyle=val_linestyle, label="验证损失")
    # 绘制坐标轴和图例
    plt.ylabel("损失", fontsize=fontsize)
    plt.xlabel("轮次", fontsize=fontsize)
    plt.legend(loc=loss_legend_loc, fontsize='x-large')
    # 绘制评价准确率变化曲线
    if len(runner.val_scores) > 0:
        plt.subplot(1, 2, 2)
        plt.plot(val_steps, runner.val_scores,
                 color=val_color, linestyle=val_linestyle, label="验证准确率")
        # 绘制坐标轴和图例
        plt.ylabel("准确率", fontsize=fontsize)
        plt.xlabel("轮次", fontsize=fontsize)
        plt.legend(loc=acc_legend_loc, fontsize='x-large')
    plt.savefig(fig_name)
    plt.show()

Runner类

import torch


# 更改主要是加入了小批量训练相关的内容
class RunnerV3(object):
    def __init__(self, model, optimizer, metric, loss_fn, **kwargs):
        self.model = model
        self.optimizer = optimizer
        self.loss_fn = loss_fn
        self.metric = metric
        # 记录训练过程中的评价指标变化情况
        self.val_scores = []
        # 记录训练过程中的损失函数变化情况
        self.train_epoch_losses = []  # 一个epoch记录一次loss
        self.train_step_losses = []  # 一个step记录一次loss
        self.val_losses = []
        # 记录全局最优指标
        self.best_score = 0

    def train(self, train_loader, val_loader=None, **kwargs):
        # 将模型切换为训练模式
        self.model.train()
        num_epochs = kwargs.get("num_epochs", 0)  # 传入训练轮数，如果没有传入值则默认为0
        log_steps = kwargs.get("log_steps", 100)  # 传入log打印频率，如果没有传入值则默认为100
        save_dir = kwargs.get("save_dir", None)  # 传入模型保存路径
        # 评价频率
        eval_steps = kwargs.get("eval_steps", 0)
        custom_print_log = kwargs.get("custom_print_log", None)  # log打印函数，如果没有传入则默认为None
        # 训练总的步数
        num_training_steps = num_epochs * len(train_loader)
        if eval_steps:
            if self.metric is None:
                raise RuntimeError('Error: Metric can not be None!')
            if val_loader is None:
                raise RuntimeError('Error: dev_loader can not be None!')
        # 运行的step数目
        global_step = 0
        # 进行num_epochs轮训练
        for epoch in range(num_epochs):
            # 用于统计训练集的损失
            total_loss = 0
            for step, data in enumerate(train_loader):
                X, y = data
                # 获取模型预测
                logits = self.model(X)
                loss = self.loss_fn(logits, y)  # 默认求mean
                total_loss += loss
                # 训练过程中，每个step的loss进行保存
                self.train_step_losses.append((global_step, loss.item()))
                if log_steps and global_step % log_steps == 0:
                    print(
                        f"[Train] epoch: {epoch}/{num_epochs}, step: {global_step}/{num_training_steps}, loss: {loss.item():.5f}")
                # 梯度反向传播，计算每个参数的梯度值
                loss.backward()
                if custom_print_log:
                    custom_print_log(self)
                # 小批量梯度下降进行参数更新
                self.optimizer.step()
                # 梯度归零
                self.optimizer.zero_grad()
                # 判断是否需要评价
                if eval_steps > 0 and global_step > 0 and \
                        (global_step % eval_steps == 0 or global_step == (num_training_steps - 1)):
                    val_score, val_loss = self.evaluate(val_loader, global_step=global_step)
                    print(f"[Evaluate]  验证准确率: {val_score:.5f}, 验证损失: {val_loss:.5f}")
                    # 将模型切换为训练模式
                    self.model.train()
                    # 如果当前指标为最优指标，保存该模型
                    if val_score > self.best_score:
                        self.save_model(save_dir)
                        print(
                            f"[Evaluate] 指标为最优，保存模型：{self.best_score:.5f} --> {val_score:.5f}")
                        self.best_score = val_score
                global_step += 1
                # 当前epoch 训练loss累计值
            trn_loss = (total_loss / len(train_loader)).item()
            # epoch粒度的训练loss保存
            self.train_epoch_losses.append(trn_loss)
        print("[Train] 训练结束!")

    # 模型评估阶段，使用'torch.no_grad()'控制不计算和存储梯度
    @torch.no_grad()
    def evaluate(self, val_loader, **kwargs):
        assert self.metric is not None
        # 将模型设置为评估模式
        self.model.eval()
        global_step = kwargs.get("global_step", -1)
        # 用于统计训练集的损失
        total_loss = 0
        # 重置评价
        self.metric.reset()
        # 遍历验证集每个批次
        for batch_id, data in enumerate(val_loader):
            X, y = data
            # 计算模型输出
            logits = self.model(X)
            # 计算损失函数
            loss = self.loss_fn(logits, y).item()
            # 累积损失
            total_loss += loss
            # 累积评价
            self.metric.update(logits, y)
        dev_loss = (total_loss / len(val_loader))
        dev_score = self.metric.accumulate()
        # 记录验证集loss
        if global_step != -1:
            self.val_losses.append((global_step, dev_loss))
            self.val_scores.append(dev_score)
        return dev_score, dev_loss

    # 模型评估阶段，使用'torch.no_grad()'控制不计算和存储梯度
    @torch.no_grad()
    def predict(self, x, **kwargs):
        # 将模型设置为评估模式
        self.model.eval()
        # 运行模型前向计算，得到预测值
        logits = self.model(x)
        return logits

    def save_model(self, save_path):
        torch.save(self.model.state_dict(), save_path)

    def load_model(self, model_path):
        model_state_dict = torch.load(model_path)
        self.model.load_state_dict(model_state_dict)

Accuracy函数

import torch


class Accuracy(object):
    def __init__(self, is_logist=True):
        # 用于统计正确的样本个数
        self.num_correct = 0
        # 用于统计样本的总数
        self.num_count = 0
        self.is_logist = is_logist

    def update(self, outputs, labels):
        # 判断是二分类任务还是多分类任务，shape[1]=1时为二分类任务，shape[1]>1时为多分类任务
        if outputs.shape[1] == 1:  # 二分类
            outputs = torch.squeeze(outputs, dim=-1)
            if self.is_logist:
                # logist判断是否大于0
                preds = (outputs >= 0).to(torch.float32)
            else:
                # 如果不是logist，判断每个概率值是否大于0.5，当大于0.5时，类别为1，否则类别为0
                preds = (outputs >= 0.5).to(torch.float32)
        else:
            # 多分类时，使用'torch.argmax'计算最大元素索引作为类别
            preds = torch.argmax(outputs, dim=1)

        # 获取本批数据中预测正确的样本个数
        labels = torch.squeeze(labels, dim=-1)
        batch_correct = torch.sum((preds == labels).float()).item()  # 直接转换为浮点数
        batch_count = len(labels)
        # 更新num_correct 和 num_count
        self.num_correct += batch_correct
        self.num_count += batch_count

    def accumulate(self):
        # 使用累计的数据，计算总的指标
        if self.num_count == 0:
            return 0
        return self.num_correct / self.num_count

    def reset(self):
        # 重置正确的数目和总数
        self.num_correct = 0
        self.num_count = 0

    def name(self):
        return "Accuracy"

预测手写数字小Demo

import torch
from torchvision import transforms
from PIL import Image
import torch.nn.functional as F
import torch.nn as nn


# -----------------------------------------------------构建网络----------------------------------------------------------
class PyTorch_LeNet(nn.Module):
    def __init__(self, in_channels, num_classes=10):
        super(PyTorch_LeNet, self).__init__()
        # 卷积层：输出通道数为6，卷积核大小为5*5
        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=6, kernel_size=5)
        # 汇聚层：汇聚窗口为2*2，步长为2
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        # 卷积层：输入通道数为6，输出通道数为16，卷积核大小为5*5
        self.conv3 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5)
        # 汇聚层：汇聚窗口为2*2，步长为2
        self.pool4 = nn.AvgPool2d(kernel_size=2, stride=2)
        # 卷积层：输入通道数为16，输出通道数为120，卷积核大小为5*5
        self.conv5 = nn.Conv2d(in_channels=16, out_channels=120, kernel_size=5)
        # 全连接层：输入神经元为120，输出神经元为84
        self.linear6 = nn.Linear(in_features=120, out_features=84)
        # 全连接层：输入神经元为84，输出神经元为类别数
        self.linear7 = nn.Linear(in_features=84, out_features=num_classes)

    def forward(self, x):
        # C1：卷积层+激活函数
        output = F.relu(self.conv1(x))
        # S2：汇聚层
        output = self.pool2(output)
        # C3：卷积层+激活函数
        output = F.relu(self.conv3(output))
        # S4：汇聚层
        output = self.pool4(output)
        # C5：卷积层+激活函数
        output = F.relu(self.conv5(output))
        # 输入层将数据拉平[B,C,H,W] -> [B,CxHxW]
        output = torch.squeeze(output, dim=3)
        output = torch.squeeze(output, dim=2)
        # F6：全连接层
        output = F.relu(self.linear6(output))
        # F7：全连接层
        output = self.linear7(output)
        return output


# 定义模型结构（确保与保存的权重一致）
model = PyTorch_LeNet(in_channels=1, num_classes=10)
model.load_state_dict(
    torch.load('C:/Users/chong/Desktop/20221205013/Python/Py/Deep Learning/Train_logs/best_model.pth'))
model.eval()  # 设置模型为评估模式
# 定义图像预处理步骤
transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=1),  # 转换为单通道灰度图
    transforms.Resize((32, 32)),  # 调整图像尺寸以匹配模型输入
    transforms.ToTensor(),  # 转换为张量
    transforms.Normalize((0.5,), (0.5,))  # 可根据数据集均值和标准差进行归一化
])

# 加载和预处理图像
image_path = 'three.png'  # 替换为你的图片路径
image = Image.open(image_path)
image = transform(image).unsqueeze(0)  # 增加批量维度，使形状为 [1, 1, 32, 32]
# 将图像传入模型进行预测
with torch.no_grad():  # 不计算梯度，加快推理速度并节省内存
    output = model(image)
    predicted_label = torch.argmax(output, dim=1).item()

print(f'Predicted label: {predicted_label}')