Pytorch CNN网络MNIST数据集图片识别

Pytorch CNN网络MNIST数据集图片识别
1、配置库和配置参数
2、加载MNIST数据集
3、创建CNN模型
4、训练模型
5、模型评估

#1、配置库和配置参数

import torch
from torch import nn, optim
import torch.nn.functional as F
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
import time

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

#配置参数
torch.manual_seed(1)    #设置随机种子，确保结果可重复性
batch_size = 128    #批处理大小
learning_rate = 1e-2    #学习率
num_epoches = 2    #训练次数

#2、加载MNIST数据集
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize([0.5], [0.5])]
)

#2、加载MNIST数据集
train_dataset = datasets.MNIST(root=r'D:\file\PythonProjict\dataset', train=True, transform=transform, download=True)   #保存位置，训练集数据，转换为Tensor，下载数据
test_dataset = datasets.MNIST(root=r'D:\file\PythonProjict\dataset', train=False, transform=transform)

#训练数据的加载方式：每次从train_dataset中随机（shuffle=True）选择batch_size个样本作为一个批次返回，因此所选择的数据可能会重复
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
#测试数据的加载方式：每次从test_dataset中选择batch_size个不同的样本作为一个批次返回，要覆盖到所有测试样本
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

#3、创建CNN模型

class Cnn(nn.Module):
    def __init__(self):
        super(Cnn, self).__init__()
        # 卷积层
        self.conv = nn.Sequential(
            nn.Conv2d(1, 25, 3),
            nn.BatchNorm2d(25),
            nn.ReLU(True),
            nn.MaxPool2d(2, 2),

            nn.Conv2d(25, 50, 3),
            nn.BatchNorm2d(50),
            nn.ReLU(True),
            nn.MaxPool2d(2, 2)
        )
        # 全连接层
        self.fc = nn.Sequential(
            nn.Linear(50*5*5, 1024),
            nn.ReLU(True),
            nn.Linear(1024, 128),
            nn.ReLU(True),
            nn.Linear(128, 10)
        )
    # 前向通道
    def forward(self, x):
        out = self.conv(x)
        out = out.view(x.size(0), -1)
        out = self.fc(out)
        return out

model = Cnn().to(device)

#4、训练模型

#定义损失：交叉熵损失
criterion = nn.CrossEntropyLoss()
#定义优化器：随梯度下降SGD
optimizer = optim.SGD(model.parameters(), lr=learning_rate)
#开始训练
model.train()
epoch = 0
train_loss = 0
train_accuracy = 0
for epoch in range(num_epoches):
    for img, label in train_loader:
        if torch.cuda.is_available():
            img = img.cuda()
            label = label.cuda()
        else:
            img = Variable(img)
            label = Variable(label)
        #前向传播
        out = model(img)
        loss = criterion(out, label)
        optimizer.zero_grad()
        train_loss += loss.item()*label.size(0)
        _, pred = torch.max(out, 1)
        num_correct = (pred == label).sum()
        train_accuracy += num_correct.item()
        #向后传播
        loss.backward()
        optimizer.step()
    print('the {}th trainning completed!'.format(epoch+1))
print('Train Times: {}, train_loss: {:.4}, train_accuracy:{:.4f}%'.format(epoch+1, train_loss/len(train_dataset)/num_epoches, 100*train_accuracy/len(train_dataset)/num_epoches))

# 5、模型评估
# with torch.no_grad():

test_start = time.time()
test_loss = 0
test_accuracy = 0
for epoch, (img, label) in enumerate(test_loader, 0):
    if torch.cuda.is_available():
        img = img.cuda()
        label = label.cuda()
    else:
        img = Variable(img)
        label = Variable(label)
    out = model(img)
    loss = criterion(out, label)
    test_loss += loss.data.item()*label.size(0)
    _, pred = torch.max(out, 1)
    num_correct = (pred == label).sum()
    test_accuracy += num_correct.item()
print('Test times: {}, test_loss: {:.4f}, test_accuracy: {:.4f}%'.format(epoch+1, test_loss / (len(test_dataset)), 100*test_accuracy / (len(test_dataset))))

结果