搭建线性网络对MNIST数据集进行训练、测试，并且预测图片

import torch import torchvision as tv # 提供数据集 import torchvision.transforms as transforms # 图像处理包 import torch.nn as nn import torch.nn.functional as F from torch import optim # 优化器包 from model import Net # 导入建立的网络 transform = transforms.Compose([transforms.ToTensor(), # 改变通道顺序、归一化 transforms.Normalize((0.5,),(0.5,))]) # 归一化 batch_size = 64 # 训练集 6W张图片 trainset =tv.datasets.MNIST(root= './dataset/mnist',train=True,download=True,transform=transform) trainloader = torch.utils.data.DataLoader(trainset,shuffle = True,batch_size = batch_size) # 测试集 1W张图片 testset =tv.datasets.MNIST(root= './dataset/mnist',train=False,download=True,transform=transform) testloader = torch.utils.data.DataLoader(testset,shuffle = False,batch_size = batch_size) net = Net() # 实例化网络 criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(net.parameters(),lr = 0.01,momentum=0.5) for epoch in range(5): running_loss = 0.0 for i,data in enumerate(trainloader,0): inputs,labels =data # 取出对应的数据 optimizer.zero_grad() # 梯度清零 outputs = net(inputs) # 计算预测值 loss = criterion(outputs,labels) # 计算损失值 loss.backward() # 反向传播 optimizer.step() # 梯度更新 running_loss += loss.item() # 取出loss值 if i % 300 ==299: correct = 0 total = 0 with torch.no_grad(): # 不会计算梯度 for data in testloader: # 从测试集里面取出数据 images, labels = data outputs = net(images) # 计算预测值 _, predicted = torch.max(outputs, dim=1) # 取出每一行最大值,返回值为（值，索引） total += labels.size(0) # batch_size correct += (predicted == labels).sum().item() print('[%d,%5d] train_loss: %.3f test_accuracy:%.3f' % (epoch+1,i +1,running_loss/300,correct / total)) running_loss = 0.0 print('Finished Training....') save_path = './Net.pth' torch.save(net.state_dict(),save_path)