pytorch RNN

pytorch RNN

import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
import torchvision
import matplotlib.pyplot as plt



# Hyper Parameters
EPOCH = 1
BATCH_SIZE = 64
TIME_STEP = 28      # rnn time step / image height
INPUT_SIZE = 28      # rnn input size / image width
LR = 0.01          # learning rate
DOWNLOAD_MNIST = False

train_data = torchvision.datasets.MNIST(
    root='./mnist',         #保存在文件夹中
    train=True,         #false则是测试点test
    transform=torchvision.transforms.ToTensor(),     #tensor数据的值在0-1之间 ，彩的0-255压缩到黑白的0-1
    download=DOWNLOAD_MNIST
    )
train_loader = torch.utils.data.DataLoader(dataset=train_data,batch_size=BATCH_SIZE,shuffle=True,num_workers=2)

test_data = torchvision.datasets.MNIST(
    root='./mnist/',
    train=False,
    transform=torchvision.transforms.ToTensor
    )
test_x = Variable(test_data.test_data,volatile=True).type(torch.FloatTensor)[:2000]/255    #shape
test_y = test_data.test_labels[:2000]   #取前1000个标签y


class RNN(nn.Module):
    def __init__(self):
        super(RNN, self).__init__()

        self.rnn = nn.LSTM(
            input_size=INPUT_SIZE,     #28
            hidden_size=64,     # rnn hidden unit
            num_layers=1,       # number of rnn layer   hidden layer=1
            batch_first=True,   # input & output will has batch size as 1s dimension. e.g. (batch, time_step, input_size)
        )
        self.out = nn.Linear(64, 10)

    def forward(self, x):
        # x (batch, time_step, input_size)
        # h_state (n_layers, batch, hidden_size)
        # r_out (batch, time_step, hidden_size)
        r_out, (h_n,h_c) = self.rnn(x, None)
        out = self.out(r_out[:,-1,:])# save all predictions
        return out
rnn = RNN()


optimizer = torch.optim.Adam(rnn.parameters(), lr=LR)   # optimize all cnn parameters
loss_func = nn.CrossEntropyLoss()        #不是0，1 这种，标签是7他显示的就是7


for epoch in range(EPOCH):
    for step, (x, y) in enumerate(train_loader):       #give batch data
        b_x = Variable(x.view(-1,28,28))              # reshape x to (batch, time_step, input_size)
        b_y = Variable(y)                     #view()函数可以将原来的向量先转成一维向量，然后再生成指定维数的向量。
        output = rnn(b_x)
        loss = loss_func(output, y)         # calculate loss
        optimizer.zero_grad()                   # clear gradients for this training step
        loss.backward()                         # backpropagation, compute gradients
        optimizer.step()                        # apply gradients

        if step % 50 == 0:
            test_output = rnn(test_x)
            pred_y = torch.max(test_output, 1)[1].data.squeeze()
            accuracy = sum(pred_y == test_y) / test_y.size(0)
            print('Epoch:', epoch, '| train loss: %.4f' % loss.item(),'| test accuracy: %.2f' % accuracy)

    # print 10 predictions 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')