基于python通过卷积神经网络实现手写数字识别

#加载必要的库
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms

 BATCH_SIZE:每一批处理的数据，比如数据集10000，BATCH_SIZE = 128，第一次就放1-128张图片进去训练。下一个批次就从129-256。

EPOCHS：整个数据集一共训练10次。

#定义超参数
BATCH_SIZE = 128 #每批处理的数据
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(DEVICE)
EPOCHS = 10    #训练数据集的轮次

#构建pipeline，对图像做处理
pipeline = transforms.Compose([
    transforms.ToTensor(),#将图片转化为tensor
    transforms.Normalize((0.1307,), (0.3081,))#正则化
])

 收集数据集，并做预处理。

#下载、加载数据
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)

搭建网络模型

主要就是神经网络的前向传播

#构建网络模型
class Digit(nn.Module):
    def __init__(self):
      super().__init__()
      self.conv1 = nn.Conv2d(1, 10, 5)#灰度图片的通道
      self.conv2 = nn.Conv2d(10, 20, 3)  # 灰度图片的通道
      self.fc1 = nn.Linear(20*100, 500)
      self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        input_size = x.size(0)
        x = self.conv1(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2, 2)

        x = self.conv2(x)
        x = F.relu(x)

        x = x.view(input_size, -1)

        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)

        output = F.log_softmax(x, dim = 1)#计算分类后每个数字的概率
        return output

#定义优化器，损失函数，将网络部署在设备上
model = Digit().to(DEVICE)
optimizer = optim.Adam(model.parameters())
loss_fn=torch.nn.CrossEntropyLoss()

#定义训练方法
def train_model(model, device, train_loader, optimizer, epoch):
    #模型训练
    model.train()
    for batch_index, (data, target) in enumerate(train_loader):
        #部署到DEVICE上
        data, target = data.to(device), target.to(device)
        #梯度初始化为0
        optimizer.zero_grad()
        #训练后的结果
        output = model(data)
        #计算损失
        #loss = F.cross_entropy(output, target)
        loss=loss_fn(output, target)

        #找到概率值最大的下标
        # pred = output.max(1, keepdim = True)
        #反向传播
        loss.backward()
        #参数优化
        optimizer.step()
        if batch_index % 3000 == 0:
            print("Train Epoch : {}\t Loss : {:.6f}".format(epoch, loss.item()))

#定义测试方法
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:
            #部署到device
            data, target = data.to(device), target.to(device)
            #测试数据
            output = model(data)
            #计算测试损失
            #test_loss += F.loss_entropy(output, target).item()
            test_loss += loss_fn(output, target).item()
            #找到概率值最大的下标
            pred = output.max(1, keepdim = True)[1]
            #累计正确率
            correct += pred.eq(target.view_as(pred)).sum().item()
        test_loss /= len(test_loader.dataset)
        print("Test - Average loss: {:.4f}, Accuracy : {:.3f}".format(
            test_loss, 100.0 * correct/len(test_loader.dataset)))

直接调用函数

for epoch in range(1, EPOCHS + 1):
    train_model(model, DEVICE, train_loader, optimizer, epoch)
    test_model(model, DEVICE, test_loader)