Pytorch之全连接识别MNIST数字

导入库

import torch
from torch import nn,optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import datasets,transforms

设置超参数

# 超参数
batch_size = 64
learning_rate = 1e-2
num_epoches = 20

数据预处理方法

# 数据预处理。transforms.ToTensor()将图片转换成PyTorch中处理的对象Tensor,并且进行标准化（数据在0~1之间）
# transforms.Normalize()做归一化。它进行了减均值，再除以标准差。两个参数分别是均值和标准差。
# transforms.Compose()函数则是将各种预处理的操作组合到了一起
data_tf = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize([0.5], [0.5])])

# 注意transforms.Normalize([0.5], [0.5]) 是将数据从0~1之间，通过减0.5，再除0.5，缩放至-1~1之间

数据集下载及获取

# 下载数据集，如果数据集已存在，则不下载。
train_dataset = datasets.MNIST(root='./data',train=True,transform=data_tf,download=True)
test_dataset = datasets.MNIST(root='./data',train=False,transform=data_tf)

# 获得训练数据集和测试数据集
train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True)
test_loader = DataLoader(test_dataset,batch_size=batch_size,shuffle=False)

模型建立

# 建立三层全连接线性网络层，并加入非线性函数ReLU。并且在线性层和非线性层之间增加批标准化，增加收敛数据

class Batch_Net(nn.Module):
    def __init__(self,in_dim,n_hidden_1,n_hidden_2,out_dim):
        super(Batch_Net,self).__init__()
        self.layer1 = nn.Sequential(nn.Linear(in_dim,n_hidden_1),nn.BatchNorm1d(n_hidden_1),nn.ReLU(True))
        self.layer2 = nn.Sequential(nn.Linear(n_hidden_1,n_hidden_2),nn.BatchNorm1d(n_hidden_2),nn.ReLU(True))
        self.layer3 = nn.Sequential(nn.Linear(n_hidden_2,out_dim))
    
    def forward(self,x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        return x

model = Batch_Net(28*28,300,100,10)
if torch.cuda.is_available():
    model = model.cuda()

确定损失函数和优化器

# 选择交叉熵作为损失函数，并选择随机梯度下降算法为优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),lr=learning_rate)

模型训练

# 训练模型
epoch = 0

# 选取所有训练数据集进行训练

for data in train_loader:
    # 获得img(手写图片)，label标签（手写图片对应数字）
    img, label = data
    img = img.view(img.size(0), -1)
    if torch.cuda.is_available():
        img = img.cuda()
        label = label.cuda()
    else:
        # 将数据转换成tensor类型
        img = Variable(img)
        label = Variable(label)
    
    #  向前传播，获得out结果和损失函数
    out = model(img)
    loss = criterion(out, label)
    print_loss = loss.data.item()
 
    # 反向传播，更新参数
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    epoch+=1
    if epoch%50 == 0:
        print('epoch: {}, loss: {:.4}'.format(epoch, loss.data.item()))

执行结果

image.png

模型评估

# 模型评估,调整模型为测试模式
model.eval()
# 累计损失函数和准确率
eval_loss = 0
eval_acc = 0

# 对测试集进行测试
for data in test_loader:
    # 获得img(手写图片)，label标签（手写图片对应数字）
    img, label = data
    img = img.view(img.size(0), -1)
    if torch.cuda.is_available():
        img = img.cuda()
        label = label.cuda()
    
    #  向前传播，获得out结果和损失函数
    out = model(img)
    loss = criterion(out, label)
    
    # 损失函数乘标签大小累计
    eval_loss += loss.data.item()*label.size(0)
    # 在10维数据中，获得最大的预测值（即预测数）
    _, pred = torch.max(out, 1)
    # 判断是否与真实结果相同
    num_correct = (pred == label).sum()
    
    # 累计真实结果
    eval_acc += num_correct.item()

# 输出评估结果    
print('Test Loss: {:.6f}, Acc: {:.6f}'.format(
    eval_loss / (len(test_dataset)),
    eval_acc / (len(test_dataset))
))

执行结果

image.png

参考：https://blog.csdn.net/out_of_memory_error/article/details/81414986#