基于pytorch的Mnist数据集进行分类(CNN，LSTM)

绝地灬酷狼 2021-12-12 04:45 847阅读 0赞

本项目是以pytorch为框架进行mnist图像分类任务：  
**CNN：**

#coding = utf-8
    import torch
    import torch.nn as nn
    from torch.autograd import Variable
    import torch.utils.data as Data
    import torchvision
    import matplotlib.pyplot as plt
    from torch import optim
    
    #parameters
    epochs = 2
    batch_size = 100
    lr = 0.01
    download_mnist = False
    
    train_data = torchvision.datasets.MNIST(
        root='./minst_data',
        train=True,
        transform = torchvision.transforms.ToTensor(),
        download=download_mnist
    )
    
    #plot one example
    # print(train_data.data.size())
    # print(train_data.targets.size())
    # plt.imshow(train_data.data[0].numpy(),cmap = 'gray')
    # plt.title('%i'%train_data.targets[0])
    # plt.show()
    
    train_loader =Data.DataLoader(dataset=train_data,batch_size=batch_size,shuffle=True)
    
    test_data = torchvision.datasets.MNIST(root='./minst_data',train=False)
    test_x = Variable(torch.unsqueeze(test_data.data,dim=1)).type(torch.FloatTensor)[:2000]/255.
    test_y = test_data.targets[:2000]
    
    class CNN(nn.Module):
        def __init__(self):
            super(CNN,self).__init__()
            self.conv1 = nn.Sequential(
                nn.Conv2d(    #(1,28,28)1是channel的维度，28*28为图片的长宽
                    in_channels=1,#图片的层数，RGB=3,灰度=1
                    out_channels=16,#filter的个数
                    kernel_size=5,#filter的长宽
                    stride=1,#每隔多少个移动
                    padding=2,#图片补0.if stride = 1,padding =(kernel_size-1)/2 = (5-1)/2 = 2
                ),   #-->（16，28，28）
                nn.ReLU(),#-->（16，28，28）
                nn.MaxPool2d(kernel_size=2),#可以看成2*2的filter    #-->（16，14，14）kernel_size = 2,减小一半
            )
            self.conv2 = nn.Sequential(#-->（16，14，14）
                nn.Conv2d(16,32,5,1,2),#-->（32，14，14）
                nn.ReLU(),#-->（32，14，14）
                nn.MaxPool2d(2)#-->（32，7，7）
    
            )
            self.out = nn.Linear(32*7*7,10)#输出为10类
    
        def forward(self, x):
            x = self.conv1(x)
            x = self.conv2(x)   #-->(batch,32,7,7)
            x = x.view(x.size(0),-1) #-->(batch,32*7*7)
            output = self.out(x)
            return output
    
    cnn = CNN()
    print(cnn)
    
    opimizer = optim.Adam(cnn.parameters(),lr=lr)
    
    loss_func = nn.CrossEntropyLoss()
    
    #train and test
    for epoch in range(epochs):
        for step,(x,y) in enumerate(train_loader):
            b_x = Variable(x)
            b_y = Variable(y)
    
            out = cnn(b_x)
            loss = loss_func(out,b_y)
            opimizer.zero_grad()
            loss.backward()
            opimizer.step()
    
            if step % 50 == 0:
                test_out = cnn(test_x)
                pred_y = torch.max(test_out,1)[1].data.squeeze()
                accuracy = sum(pred_y == test_y)/ test_y.size
                print('Epoch:',epoch,'|train loss:'+str(loss.item()),'|test accuracy:'+str(accuracy))

**LSTM:**

#coding = utf-8
    
    
    import torch
    import torch.nn as nn
    from torch.autograd import Variable
    import torch.utils.data as Data
    import torchvision.datasets as dsets
    import torchvision.transforms as transforms
    import matplotlib.pyplot as plt
    from torch import optim
    
    #parameters
    epochs = 2
    batch_size = 64
    time_step = 28  #time step/image height
    input_size = 28 #input_size / image width
    lr = 0.01
    download_minst = False
    
    train_data = dsets.MNIST(
        root='./minst_data',
        train=True,
        transform = transforms.ToTensor(),
        download=download_minst
    )
    
    train_loader =Data.DataLoader(dataset=train_data,batch_size=batch_size,shuffle=True)
    
    test_data = dsets.MNIST(root='./minst_data',train=False,transform=transforms.ToTensor())
    test_x = Variable(test_data.data).type(torch.FloatTensor)[:2000]/255.
    test_y = test_data.targets.numpy().squeeze()[:2000]
    
    class Rnn(nn.Module):
        def __init__(self):
            super(Rnn,self).__init__()
    
            self.rnn = nn.LSTM(
                input_size=input_size,
                hidden_size=64,
                num_layers=2,
                batch_first=True,  #（time_step,batch,input）时是Ture
            )
            self.out = nn.Linear(64,10)
    
        def forward(self, x):
            r_out,(h_n,h_c) = self.rnn(x,None)  # x (batch,time_step,input_size)
            out = self.out(r_out[:,-1,:]) #(batch,time_step,input)
            return out
    
    rnn = Rnn()
    print(rnn)
    
    optimizer = optim.Adam(rnn.parameters(),lr=lr)
    
    loss_func = nn.CrossEntropyLoss()
    
    for epoch in range(epochs):
        for step,(x,y) in enumerate(train_loader):
            b_x = Variable(x.view(-1,28,28))
            b_y = Variable(y)
            output = rnn(b_x)
            loss = loss_func(output,b_y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    
            if step % 50 == 0:
                test_out = rnn(test_x)#(sample,time_step,input_size)
                pred_y = torch.max(test_out,1)[1].data.numpy().squeeze()
                accuracy = sum(pred_y == test_y) / test_y.size
                print('Epoch:', epoch, '|train loss:' + str(loss.item()), '|test accuracy:' + str(accuracy))
    
    #print 10 prediction from test data
    test_output = rnn(test_x[:10].view(-1,28,28))
    pred_y = torch.max(test_output,1)[1].data.numpy().squeeze()
    print(pred_y,'prediction number')
    print(test_y[:10],'real number')

代码如上所示；有问题请留言，欢迎大家积极讨论！！！