pytorch 训练MNIST数据集

使用pytorch1.3 训练MNIST数据集

本次博客的要点在于，分析主要代码的功能作用，便于理解。
参考：(https://blog.csdn.net/weicao1990/article/details/98754647)

import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn import functional as F
from torchvision import datasets, transforms

#设置批次，学习率，训练的次数
batch_size = 500
learning_rate = 0.01
epochs = 10

#加载MNIST训练集，测试集
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('../data', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ])),
    batch_size=batch_size, shuffle=True)
#测试集
test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('../data', train=False, transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])),
    batch_size=batch_size, shuffle=True)

#设置权重参数w及偏置参数b， 并设置梯度计算
w1,b1 = torch.randn(200,784,requires_grad=True),\
    torch.zeros(200,requires_grad=True)
w2,b2 = torch.randn(200,200,requires_grad=True),\
    torch.zeros(200,requires_grad=True)
w3,b3 = torch.randn(10,200,requires_grad=True),\
    torch.zeros(10,requires_grad=True)

#对w初始化
torch.nn.init.kaiming_normal_(w1)
torch.nn.init.kaiming_normal_(w2)
torch.nn.init.kaiming_normal_(w3)

#前向传播的架构
def forward(x):
    x = x@w1.t() + b1
    x = F.relu(x)
    x = x@w2.t() + b2
    x = F.relu(x)
    x = x@w3.t() + b3
    x = F.relu(x)
    return x

#设置优化器及损失函数
optimizer = optim.SGD([w1,b1,w2,b2,w3,b3],lr=learning_rate)
criteon = nn.CrossEntropyLoss()

#进行训练
for epoch in range(epochs):
	#batch_idx表示载入批次数据的次数，target表示data对应的目标值
    for batch_idx,(data,target) in enumerate(train_loader):
        data = data.view(-1,28*28)
        #将数据导入到前项传播进行计算，并得到损失函数
        logits = forward(data)
        loss = criteon(logits,target)
        #优化器梯度初始化
        optimizer.zero_grad()
        #对损失函数进行反向传播，计算梯度
        loss.backward()
        #更新梯度
        optimizer.step()
		#每加载一百次输出一次信息
        if batch_idx % 100 == 0:
            print('Train Epoch:{},{}, [{}/{} ({:.0f}%)]\tLoss:{:.6f}'\
                  .format(epoch,batch_idx, batch_idx*len(data),len(train_loader.dataset),\
                          100.*batch_idx / len(train_loader),loss.item()))

    test_loss = 0
    correct = 0
    for data,target in test_loader:
        data = data.view(-1,28*28)
        logits = forward(data)
        #item()表示转化为numpy的格式
        test_loss += criteon(logits,target).item()
		# 得到的预测值输出是一个10个分量的概率,在第2个维度上取max
		# logits.data是一个shape=[batch_size,10]的Tensor
		# 注意Tensor.max(dim=1)是在这个Tensor的1号维度上求最大值
		# 得到一个含有两个元素的元组,这两个元素都是shape=[batch_size]的Tensor
		# 第一个Tensor里面存的都是最大值的值,第二个Tensor里面存的是对应的索引
		# 这里要取索引,所以取了这个tuple的第二个元素
		# print(type(logits.data), logits.data.shape,type(logits.data.max(dim=1)))
		# 等同于pred = logits.argmax(dim=1)
        pred = logits.data.max(1)[1]
        #累加与与目标值相等的个数，作为准确的数量
        correct += pred.eq(target.data).sum()

    test_loss /= len(test_loader.dataset)
    print('\nTest set: Average loss:{:.4f}, Accuracy:{}/{} ({:.0f}%)\n)'.format\
              (test_loss,correct,len(test_loader.dataset),100.*\
               correct / len(test_loader.dataset)))