1.使用RNN做MNIST分类

第一次用LSTM,从简单做起吧~~

注意事项：

batch_first=True 意味着输入的格式为（batch_size,time_step,input_size),False 意味着输入的格式为（time_step,batch_size,input_size)
取r_out[:,-1,:]，即取时间步最后一步的结果，相当于LSTM把一张图片全部扫描完后的返回的状态向量（此时的维度变为（64，64），前面的64是batch_size,后面的64是隐藏层的神经元个数）

 import torch

 from torch.autograd import Variable

 from torchvision import datasets,transforms

 #超参数

 EPOCH=1

 BATCH_SIZE=64

 TIME_STEP=28#run time step/image height

 INPUT_SIZE=28#run input size/image width

 LR=0.01

 DOWNLOAD_MNIST=True

 train_data=datasets.MNIST(root='./mnist',train=True,transform=transforms.ToTensor(),download=DOWNLOAD_MNIST)

 train_loader=torch.utils.data.DataLoader(dataset=train_data,batch_size=BATCH_SIZE,shuffle=True)

 test_data=datasets.MNIST(root='./mnist',train=False,transform=transforms.ToTensor(),download=DOWNLOAD_MNIST)

 test_loader=torch.utils.data.DataLoader(dataset=test_data,batch_size=BATCH_SIZE,shuffle=True)

 class RNN(torch.nn.Module):

     def __init__(self):

         super(RNN,self).__init__()

         self.rnn=torch.nn.LSTM(

             input_size=INPUT_SIZE,

             hidden_size=64,

             num_layers=1,

             batch_first=True,

         )

         self.out=torch.nn.Linear(64,10)

     def forward(self, x):

         r_out,(h_n,h_c)=self.rnn(x,None)#[64,28,64]

         out=self.out(r_out[:,-1,:])#[64,10]

         return out

 #time_step,batch,input  batch_first=False,

 rnn=RNN()

 print(rnn)

 optimizer=torch.optim.Adam(rnn.parameters(),lr=LR)

 loss_func=torch.nn.CrossEntropyLoss()

 for epoch in range(EPOCH):

     for step,(x,y) in enumerate(train_loader):

         b_x=Variable(x.view(-1,28,28))#reshape x to (batch,time_step.input_size)

         b_y=Variable(y).squeeze()

         output=rnn(b_x)

         loss=loss_func(output,b_y)

         optimizer.zero_grad()

         loss.backward()

         optimizer.step()

         if step %50==0:

             for test_x,test_y in test_loader:

                 test_output=rnn(test_x.view(-1,28,28))

                 pred_y=torch.max(test_output,1)[1].data.numpy().squeeze()

                 test_y=test_y.numpy()

                 acc=sum(pred_y==test_y)/test_y.size

                 print(acc)

秒客网

1.使用RNN做MNIST分类

相关文章