PyTorch手写字体识别MNIST

手写字体识别MNIST

1.准备工作

可以看这个老师的视频进行学习，讲解的非常仔细：视频学习

2.项目代码

2.1 导入模块

# 1.加载相关库
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets,transforms

2.2 定义超参数
每批处理数据为16，训练轮数为10

# 2.定义超参数
BATCH_SIZE=16  #每批处理的数据
DEVICE=torch.device("cuda" if torch.cuda.is_available() else "cpu")  #是否用GPU还是CPU训练
EPOCHS=10  #训练数据集的轮次

2.3 构建pipeline，对图像做处理

# 3.构建pipeline，对图像做处理
pipeline=transforms.Compose([
    transforms.ToTensor(),#将图片转换成tensor
    transforms.Normalize((0.1307,),(0.3081))  #正则化：降低模型复杂度
])

2.4 下载、加载数据集（需要联网）
如图所示：

# 4.下载、加载数据集
from torch.utils.data import DataLoader

# 下载数据集
train_set=datasets.MNIST("data",train=True,download=True,transform=pipeline)

test_set=datasets.MNIST("data",train=False,download=True,transform=pipeline)

# 加载数据
train_loader=DataLoader(train_set,batch_size=BATCH_SIZE,shuffle=True)

test_loader=DataLoader(test_set,batch_size=BATCH_SIZE,shuffle=True)

2.5 构建网络模型（两层）

# 5.构建网络模型
class Digit(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1=nn.Conv2d(1,10,5)  #1：灰度图片的通道 10：输出通道  5：kernel
        self.conv2=nn.Conv2d(10,20,3)  #10:输入通道  20：输出通道  3：kernel
        self.fc1=nn.Linear(20*10*10,500)   #20*10*10：输入通道 500：输出通道
        self.fc2=nn.Linear(500,10)  #输入通道 10:输出通道
        
    def forward(self,x):
        input_size=x.size(0)  #atch_size*1*28*28 只取batch_size
        x=self.conv1(x)  #输入：batch*1*28*28  输出：batch*10*24*24 （28-5+1=24）
        x=F.relu(x)  #保持shape不变，输出：batch*10*24*24
        x=F.max_pool2d(x,2,2)  #输入：batch*10*24*24  输出：batch*10*12*12
        
        x=self.conv2(x)  #输入：batch*10*12*12 输出：batch*20*10*10 （12-3+1=10）
        x=F.relu(x)
        
        x=x.view(input_size,-1)  #拉平，-1 自动计算维度  20*10*10=2000
        
        x=self.fc1(x)  #输入：batch*2000  输出：batch*500
        x=F.relu(x)  #保持shape不变
        
        x=self.fc2(x)  #输入：batch*500 输出：batch*10
        
        output=F.log_softmax(x,dim=1)  #计算分类后，每个数字的概率值,那个数字的概率最大那么就输出数字几
        
        return  output

2.6 定义优化器，训练方法和测试方法

# 6.定义优化器
model=Digit().to(DEVICE)

optimizer=optim.Adam(model.parameters())

# 7.定义训练方法
def train_model(model,device,train_loader,optimizer,epoch):
    model.train()  #模型训练
    for batch_index,(data,target) in enumerate(train_loader):
        data,target=data.to(device),target.to(device)  #部署到DEVICE上去
        optimizer.zero_grad()  #梯度初始化为0
        output=model(data) #训练后的结果
        loss=F.cross_entropy(output,target)  #计算损失
#         pred=output.max(1,keepdim=True)  #pred=output.argmax(dim=1)  #找到概率值最大的下标
        loss.backward()  #反向传播
        optimizer.step()  #参数优化
        if batch_index%3000==0:
            print("Train Epoch:{} \t Loss:{:.6f}".format(epoch,loss.item()))

# 8.定义测试方法
def test_model(model,device,test_loader):
    model.eval()  #模型验证
    correct=0.0  #正确率
    test_loss=0.0  #测试损失
    with torch.no_grad():  #不会计算梯度，也不会进行反向传播
        for data,target in test_loader:
            data,target=data.to(device),target.to(device)  #部署到DEVICE上
            output=model(data) #测试数据
            test_loss+=F.cross_entropy(output,target).item()  #计算测试损失
            pred=output.max(1,keepdim=True)[1]  #[0]值  [1]索引  找到概率值最大的下标
#           pred=torch.max(output,dimm=1)
#           pred=output.argmax(dim=1)
            correct+=pred.eq(target.view_as(pred)).sum().item()  #累计正确的值
        test_loss/=len(test_loader.dataset)
        print("Test -- Average loss:{:.4f},Accuracy:{:.3f}\n".format(test_loss,100.0*correct/len(test_loader.dataset)))

2.7 调用以上方法

# 9.调用方法
for epoch in range(1,EPOCHS+1):
    train_model(model,DEVICE,train_loader,optimizer,epoch)
    test_model(model,DEVICE,test_loader)

3 运行结果

从结果中我们看到，在Epoch 8中准确率达到了99.090，说明这个模型不错。当然，我们也可以再改进一下网络模型，还请多多指教，一起学习哈！！！