pytorch实现LeNet网络对MNIST数据集分类

import torch
from torch import nn
from torch.nn import functional as F
from torch import optim
from keras.utils import to_categorical
import numpy as np

path="F:\mnist.npz"
f = np.load(path)
train_X, train_y = f['x_train'], f['y_train']
test_X, test_y = f['x_test'], f['y_test']
f.close()

train_X = train_X.reshape(-1, 28, 28, 1)
train_X = train_X.astype('float32')
train_X /= 255
train_y = to_categorical(train_y, 10)

# 创建网络

class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        self.conv1=nn.Conv2d(1,32,5,stride=1,padding=0)
        self.relu1=nn.ReLU()
        self.pool1=nn.MaxPool2d(kernel_size=2,stride=2,padding=0)
        self.conv2=nn.Conv2d(32,16,3,stride=1,padding=0)
        self.relu2=nn.ReLU()
        self.pool2=nn.MaxPool2d(kernel_size=2,stride=2,padding=0)
        self.fc1 = nn.Linear(400, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x=self.conv1(x)
        x=self.relu1(x)
        x=self.pool1(x)
        x=self.conv2(x)
        x=self.relu2(x)
        x=self.pool2(x)
        x=x.view(-1,400)
        x=self.fc1(x)
        x=F.relu(x)
        x=self.fc2(x)
        x=F.relu(x)
        x=self.fc3(x)
        return x


net = Net()

optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)

train_loss = []
precision=0

for epoch in range(4):

    for i in range(600):
        x=train_X[i*100:i*100+100]
        y=train_y[i*100:i*100+100]
        x = x.reshape(-1,1,28,28)
        x = torch.from_numpy(x)        #(batch_size,input_feature_shape)
        y = torch.from_numpy(y)        #(batch_size,label_onehot_shape)

        out = net(x)

        loss = F.mse_loss(out, y)         # 计算两者的误差
        optimizer.zero_grad()             # 清空上一步的残余更新参数值
        loss.backward()                   # 误差反向传播, 计算参数更新值
        optimizer.step()                  # 将参数更新值施加到 net 的 parameters 上
        train_loss.append(loss.item())

        if i % 10 == 0:
            print(epoch, i, np.mean(train_loss))
            train_loss=[]







total_correct = 0
for i in range(10000):
    x = train_X[i]
    y = train_y[i]
    x = torch.from_numpy(x)
    y = torch.from_numpy(y)
    x=x.view(1,1,28,28)
    y=y.view(1,10)
    out = net(x)
    pred = out.argmax(dim=1)  # 返回值最大的索引
    label = y.argmax(dim=1)
    correct = pred.eq(label).sum().float().item()  # 这个batch中正确的数量
    total_correct += correct

acc = total_correct / 10000.0
print('test acc:', acc)

训练4轮，结果如下