pytorch入门（六）—— 循环神经网络

循环神经网络（RNN）

import torch as tr
import torch.nn as nn
import torchvision.datasets as dsets
import torchvision.transforms as tf
import torch.utils.data as Data
from torch.autograd import Variable

sequence_length = 28
input_size = 28
hidden_size = 128
num_layers = 2
num_classes = 10
batch_size = 100
num_epochs = 2
learning_rate = 0.01

train_dataset = dsets.MNIST(
        root='./data/',
        train=True,
        transform=tf.ToTensor(),
        download=True
    )

train_loader = Data.DataLoader(dataset=train_dataset,
        batch_size=batch_size,
        shuffle=True)

test_dataset = dsets.MNIST(
        root='./data/',
        train=False,
        transform=tf.ToTensor()
    )

test_loader = Data.DataLoader(dataset=test_dataset,
        batch_size=batch_size,
        shuffle=False)


class RNN(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, num_classes):
        super(RNN, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        h0 = Variable(tr.zeros(self.num_layers, x.size(0), self.hidden_size)).cuda()
        c0 = Variable(tr.zeros(self.num_layers, x.size(0), self.hidden_size)).cuda()
        out, _ = self.lstm(x, (h0, c0))
        out = self.fc(out[:, -1, :])
        return out

rnn = RNN(input_size, hidden_size, num_layers, num_classes)
rnn.cuda()
print(rnn)
optimizer = tr.optim.Adam(rnn.parameters(), lr=learning_rate)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = Variable(images.view(-1, sequence_length, input_size)).cuda()
        labels = Variable(labels).cuda()

        optimizer.zero_grad()
        outputs = rnn(images)

        loss = loss_fn(outputs, labels)
        loss.backward()
        optimizer.step()

        if (i + 1) % 100 == 0:
            print('Epoch [%d/%d], Step [%d/%d], Loss: %.4f' % 
                    (epoch + 1, num_epochs, i + 1, len(train_dataset) // batch_size, loss.item()))

correct = 0
total = 0

for images, labels in test_loader:
    images = Variable(images.view(-1, sequence_length, input_size)).cuda()
    outputs = rnn(images)
    _, predicted = tr.max(outputs.data, 1)
    total += labels.size(0)
    correct += (predicted.cpu() == labels).sum()

print('Test Accuracy of the model on the 10000 test images: %d %%' %
        (100 * correct / total))

tr.save(rnn, 'RNN.pkl')

双向循环神经网络（BRNN）

import torch as tr
import torch.nn as nn
import torchvision.datasets as dsets
import torchvision.transforms as tf
import torch.utils.data as Data
from torch.autograd import Variable

sequence_length = 28
input_size = 28
hidden_size = 128
num_layers = 2
num_classes = 10
batch_size = 100
num_epochs = 2
learning_rate = 0.01

train_dataset = dsets.MNIST(
        root='./data/',
        train=True,
        transform=tf.ToTensor(),
        download=True
    )

train_loader = Data.DataLoader(dataset=train_dataset,
        batch_size=batch_size,
        shuffle=True)

test_dataset = dsets.MNIST(
        root='./data/',
        train=False,
        transform=tf.ToTensor()
    )

test_loader = Data.DataLoader(dataset=test_dataset,
        batch_size=batch_size,
        shuffle=False)


class BiRNN(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, num_classes):
        super(BiRNN, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, 
                batch_first=True, bidirectional=True)
        self.fc = nn.Linear(hidden_size * 2, num_classes)

    def forward(self, x):
        h0 = Variable(tr.zeros(self.num_layers * 2, x.size(0), self.hidden_size)).cuda()
        c0 = Variable(tr.zeros(self.num_layers * 2, x.size(0), self.hidden_size)).cuda()
        out, _ = self.lstm(x, (h0, c0))
        out = self.fc(out[:, -1, :])
        return out

brnn = BiRNN(input_size, hidden_size, num_layers, num_classes)
brnn.cuda()
#print(brnn)
optimizer = tr.optim.Adam(brnn.parameters(), lr=learning_rate)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = Variable(images.view(-1, sequence_length, input_size)).cuda()
        labels = Variable(labels).cuda()

        optimizer.zero_grad()
        outputs = brnn(images)

        loss = loss_fn(outputs, labels)
        loss.backward()
        optimizer.step()

        if (i + 1) % 100 == 0:
            print('Epoch [%d/%d], Step [%d/%d], Loss: %.4f' % 
                    (epoch + 1, num_epochs, i + 1, 
                        len(train_dataset) // batch_size, loss.item()))

correct = 0
total = 0

for images, labels in test_loader:
    images = Variable(images.view(-1, sequence_length, input_size)).cuda()
    outputs = brnn(images)
    _, predicted = tr.max(outputs.data, 1)
    total += labels.size(0)
    correct += (predicted.cpu() == labels).sum()

print('Test Accuracy of the model on the 10000 test images: %d %%' %
        (100 * correct / total))

tr.save(brnn, 'BiRNN.pkl')