轻松学习mnist手写数字数据集

1.加载必要的库

2.定义超参数

3.构建对图像处理

4.加载数据

5.构建网络模型

6.定义优化器

7.定义训练函数方法

8.测试

#加载库
import torch
import torch.nn as nn
import  torchvision
import torch.nn.functional as F
import torch.optim as optim
from  torchvision import  datasets,transforms
from torch.utils.data import DataLoader
from  torchvision import  datasets,transforms
batch_size=16   #每批处理的数据
device=torch.device("cude" if torch.cuda.is_available() else "cpu")#使用gpu还是cpu
epochs=10#训练的轮次
#构建pipeline
pipeline=transforms.Compose([transforms.ToTensor(),#将图像转换为tensor
                             transforms.Normalize((0.1307),(0.3081))])#降低模型复杂度
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(Digit, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, 5)
        self.conv2 = nn.Conv2d(10, 20, 3)
        self.fc1 = nn.Linear(20 * 10 * 10, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        input_size = x.size(0)  # batch_size
        x = self.conv1(x)  # 28-5-+1=24
        x = F.relu(x)
        x = F.max_pool2d(x, 2, 2)  # 输入：batch*10*24*24   输出：batch*10*12*12             找到最重要的部分
        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())


# 定义训练方法
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).long(), target.to(device).long()
        # 梯度初始化为0
        optimizer.zero_grad()
        # 训练后的结果
        output = model(data).long()
        # 计算损伤
        loss = F.cross_entropy(target.long(), output)
        # 找到概率最大的下标
        # pred=output.max(1,keepdim=True)#argmax(dim=1)
        # 反向传播
        loss.backwardI()
        # 参数优化
        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
    # 测试损失
    test_loss = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.cross_entropy(target, output.long()).item()
            pred = torch.max(output.long(), 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 * 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)