recurrent_neural_network

recurrent_neural_network

# Import necessary packages.
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

# Device configuration.
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)
print(torch.cuda.get_device_name())

cuda
NVIDIA GeForce GTX 1650

# Set the Hyper-parameters.
sequence_length = 28
input_size = 28
hidden_size = 128
num_layers = 2
num_classes = 10
batch_size = 100
num_epochs = 5
learning_rate = 0.01

# Load the MNIST dataset.
train_data = torchvision.datasets.MNIST(root='../../data/',
                                        train=True,
                                        transform=transforms.ToTensor(),
                                        download=True)

test_data = torchvision.datasets.MNIST(root='../../data/',
                                       train=False,
                                       transform=transforms.ToTensor())

# Define the Data Loader.
train_loader = torch.utils.data.DataLoader(dataset=train_data,
                                           batch_size=batch_size,
                                           shuffle=True)

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

# Recurrent neural network (many-to-one).
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):
        # set initial hidden and cell status.
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)

        # Forward propagate LSTM.
        out, _ = self.lstm(x, (h0, c0))  # out: tensor of shape (batch_size, seq_length, hidden_size).

        # Decode the hidden state of the last time step.
        out = self.fc(out[:, -1, :])
        return out

model = RNN(input_size, hidden_size, num_layers, num_classes).to(device)

# Loss and optimizer.
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# Train the model.
total_step = len(train_loader)
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.reshape(-1, sequence_length, input_size).to(device)
        labels = labels.to(device)

        # Forward pass.
        outputs = model(images)
        loss = criterion(outputs, labels)

        # Backward and optimize.
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Set an output counter.
        if (i+1) % 100 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
                  .format(epoch+1, num_epochs, i+1, total_step, loss.item()))

Epoch [1/5], Step [100/600], Loss: 0.2089
Epoch [1/5], Step [200/600], Loss: 0.1498
Epoch [1/5], Step [300/600], Loss: 0.1869
Epoch [1/5], Step [400/600], Loss: 0.1366
Epoch [1/5], Step [500/600], Loss: 0.0580
Epoch [1/5], Step [600/600], Loss: 0.1454
Epoch [2/5], Step [100/600], Loss: 0.1133
Epoch [2/5], Step [200/600], Loss: 0.0952
Epoch [2/5], Step [300/600], Loss: 0.0493
Epoch [2/5], Step [400/600], Loss: 0.1025
Epoch [2/5], Step [500/600], Loss: 0.1136
Epoch [2/5], Step [600/600], Loss: 0.2401
Epoch [3/5], Step [100/600], Loss: 0.1501
Epoch [3/5], Step [200/600], Loss: 0.1444
Epoch [3/5], Step [300/600], Loss: 0.0374
Epoch [3/5], Step [400/600], Loss: 0.0263
Epoch [3/5], Step [500/600], Loss: 0.0663
Epoch [3/5], Step [600/600], Loss: 0.0181
Epoch [4/5], Step [100/600], Loss: 0.0586
Epoch [4/5], Step [200/600], Loss: 0.0449
Epoch [4/5], Step [300/600], Loss: 0.0645
Epoch [4/5], Step [400/600], Loss: 0.0477
Epoch [4/5], Step [500/600], Loss: 0.1943
Epoch [4/5], Step [600/600], Loss: 0.0971
Epoch [5/5], Step [100/600], Loss: 0.0664
Epoch [5/5], Step [200/600], Loss: 0.0504
Epoch [5/5], Step [300/600], Loss: 0.0369
Epoch [5/5], Step [400/600], Loss: 0.0615
Epoch [5/5], Step [500/600], Loss: 0.0660
Epoch [5/5], Step [600/600], Loss: 0.0057

# Test the model.
model.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.reshape(-1, sequence_length, input_size).to(device)
        labels = labels.to(device)

        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

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

Test Accuracy of the model on the 10000 test images: 97.85 %

# Save the model checkpoint.
torch.save(model.state_dict(), 'model_param.ckpt')
# Alternative method to save the hole model.
# torch.save(model, 'model.ckpt')

# Load the model checkpoint.
model.load_state_dict(torch.load('model_param.ckpt'))

<All keys matched successfully>