Pytorch框架实现mnist手写库识别（与tensorflow对比）

前言
最近在学习过程中需要用到pytorch框架，简单学习了一下，写了一个简单的案例，记录一下pytorch中搭建一个识别网络基础的东西。对应一位博主写的tensorflow的识别mnist数据集，将其改为pytorch框架，也可以详细看到两个框架大体的区别。

Tensorflow版本转载来源（CSDN博主「兔八哥1024」）：https://blog.csdn.net/lzx159951/article/details/100577182

Pytorch实战mnist手写数字识别

#需要导入的包
import torch
import torch.nn as nn#用于构建网络层
import torch.optim as optim#导入优化器
from torch.utils.data import DataLoader#加载数据集的迭代器
from torchvision import datasets, transforms#用于加载mnsit数据集

#下载数据集

train_set = datasets.MNIST('./data', train=True, download=True,transform = transforms.Compose([
                  transforms.ToTensor(),
                  transforms.Normalize((0.1037,), (0.3081,))
              ]))
test_set = datasets.MNIST('./data', train=False, download=True,transform = transforms.Compose([
                  transforms.ToTensor(),
                  transforms.Normalize((0.1037,), (0.3081,))
              ]))

#构建网络（网络结构对应tensorflow的那一篇文章）

class Net(nn.Module):

    def __init__(self, num_classes=10):
        super(Net, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=5, stride=1, padding=2),
            nn.MaxPool2d(kernel_size=2,stride=2),
            nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2),
            nn.MaxPool2d(kernel_size=2,stride=2),

        )
        self.classifier = nn.Sequential(
            nn.Linear(3136, 7*7*64),
            nn.Linear(3136, num_classes),

        )

    def forward(self,x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)

        return x
net=Net()
net.cuda()#用GPU运行

#计算误差，使用adam优化器优化误差
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), 1e-2)

train_data = DataLoader(train_set, batch_size=128, shuffle=True)
test_data = DataLoader(test_set, batch_size=128, shuffle=False)



#训练过程
for epoch in range(10000):
    net.train()  ##在进行训练时加上train()，测试时加上eval()
    batch = 0

    for batch_images, batch_labels in train_data:

        average_loss = 0
        train_acc = 0

        ##在pytorch0.4之后将Variable 与tensor进行合并，所以这里不需要进行Variable封装
        if torch.cuda.is_available():
            batch_images, batch_labels = batch_images.cuda(),batch_labels.cuda()

        #前向传播
        out = net(batch_images)
        loss = criterion(out,batch_labels)


        average_loss = loss
        prediction = torch.max(out,1)[1]
        # print(prediction)

        train_correct = (prediction == batch_labels).sum()
        ##这里得到的train_correct是一个longtensor型，需要转换为float

        train_acc = (train_correct.float())  / 128

        optimizer.zero_grad() #清空梯度信息，否则在每次进行反向传播时都会累加
        loss.backward()  #loss反向传播
        optimizer.step()  ##梯度更新

        batch+=1
        print("Epoch: %d/%d || batch:%d/%d average_loss: %.3f || train_acc: %.2f"
              %(epoch, 20, batch, float(int(50000/128)), average_loss, train_acc))

# 在测试集上检验效果
net.eval()  # 将模型改为预测模式
for idx,(im1, label1) in enumerate(test_data):
    if torch.cuda.is_available():
        im, label = im1.cuda(),label1.cuda()
    out = net(im)
    loss = criterion(out, label)

    eval_loss = loss

    pred = torch.max(out,1)[1]
    num_correct = (pred == label).sum()
    acc = (num_correct.float())/ 128
    eval_acc = acc

    print('EVA_Batch:{}, Eval Loss: {:.6f}, Eval Acc: {:.6f}'
      .format(idx,eval_loss , eval_acc))

运行结果：

如果想要获取更多资源以及技术请加入我们（qq群、微信群以及公众号二维码）：