train.py

import gzip
import os
import pickle
import random

from torchvision import transforms

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim

BATCH_SIZE = 32  # mb size
EPOCHS = 30  # number of epochs
TRAIN_VAL_SPLIT = 0.85  # train/val ratio

DATA_DIR = 'data'
DATA_FILE = 'data.gzip'
MODEL_FILE = 'model.pt'

# Set of all Actions
actions_set = [[0, 0, 0],  # no action
                     [-1, 0, 0],  # left
                     [1, 0, 0],  # right
                     [0, 1, 0],  # acceleration
                     [0, 0, 1], ]  # break

# transformations for training/testing
data_transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Grayscale(1),
    transforms.Pad((12, 12, 12, 0)),
    transforms.CenterCrop(84),
    transforms.ToTensor(),
    transforms.Normalize((0,), (1,)),
])

def read_data():
    """Read the data generated by keyboard_agent.py"""
    with gzip.open(os.path.join(DATA_DIR, DATA_FILE), 'rb') as f:
        data = pickle.load(f)

    random.shuffle(data)

    states, actions, _, _, _ = map(np.array, zip(*data))

    act_classes = np.full((len(actions)), -1, dtype=int)
    for i, a in enumerate(actions_set):
        act_classes[np.all(actions == a, axis=1)] = i

    # drop unsupported actions
    states = np.array(states)
    states = states[act_classes != -1]
    act_classes = act_classes[act_classes != -1]

    for i, a in enumerate(actions_set):
        print("Actions of type {}: {}"
              .format(str(a), str(act_classes[act_classes == i].size)))

    print("Total transitions: " + str(len(act_classes)), act_classes[0], act_classes[1], act_classes[2])
    return states, act_classes


def create_datasets():

    class TensorDatasetTransforms(torch.utils.data.TensorDataset):

        def __init__(self, x, y):
            super().__init__(x, y)

        def __getitem__(self, index):
            tensor = data_transform(self.tensors[0][index])
            return (tensor,) + tuple(t[index] for t in self.tensors[1:])

    x, y = read_data()

    x = np.moveaxis(x, 3, 1)  # channel first (torch requirement)

    # train dataset
    x_train = x[:int(len(x) * TRAIN_VAL_SPLIT)]
    y_train = y[:int(len(y) * TRAIN_VAL_SPLIT)]

    train_set = TensorDatasetTransforms(
        torch.tensor(x_train),
        torch.tensor(y_train))

    train_loader = torch.utils.data.DataLoader(train_set,
                                               batch_size=BATCH_SIZE,
                                               shuffle=True,
                                               num_workers=2)

    # test dataset
    x_val, y_val = x[int(len(x_train)):], y[int(len(y_train)):]

    val_set = TensorDatasetTransforms(
        torch.tensor(x_val),
        torch.tensor(y_val))

    val_loader = torch.utils.data.DataLoader(val_set,
                                             batch_size=BATCH_SIZE,
                                             shuffle=False,
                                             num_workers=2)

    return train_loader, val_loader

def create_ex_datasets():

    class TensorDatasetTransforms(torch.utils.data.TensorDataset):

        def __init__(self, x, y):
            super().__init__(x, y)

        def __getitem__(self, index):
            tensor = data_transform(self.tensors[0][index])
            return (tensor,) + tuple(t[index] for t in self.tensors[1:])

    x, y = read_data()

    x = np.moveaxis(x, 3, 1)
    x_ex = x[:2]
    y_ex = y[:2]

    ex_set = TensorDatasetTransforms(
        torch.tensor(x_ex),
        torch.tensor(y_ex))

    ex_loader = torch.utils.data.DataLoader(ex_set)

    return ex_loader

def Net():

    class Flatten(nn.Module):

        def forward(self, x):
            return x.view(x.size()[0], -1)

    model = torch.nn.Sequential(
        torch.nn.Conv2d(1, 32, 8, 4),
        torch.nn.BatchNorm2d(32),
        torch.nn.ELU(),
        torch.nn.Dropout2d(0.5),
        torch.nn.Conv2d(32, 64, 4, 2),
        torch.nn.BatchNorm2d(64),
        torch.nn.ELU(),
        torch.nn.Dropout2d(0.5),
        torch.nn.Conv2d(64, 64, 3, 1),
        torch.nn.ELU(),
        Flatten(),
        torch.nn.BatchNorm1d(64 * 7 * 7),
        torch.nn.Dropout(),
        torch.nn.Linear(64 * 7 * 7, 120),
        torch.nn.ELU(),
        torch.nn.BatchNorm1d(120),
        torch.nn.Dropout(),
        torch.nn.Linear(120, len(actions_set)),
    )

    return model


def train(model):
    """
    Training main method
    :param model: the network
    """
    loss_function = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters())
    train_loader, val_order = create_datasets()  # read datasets

    # train
    for epoch in range(EPOCHS):
        print('Epoch {}/{}'.format(epoch + 1, EPOCHS))
        train_epoch(model, loss_function, optimizer, train_loader)
        test(model, loss_function, val_order)

        # save model
        model_path = os.path.join(DATA_DIR, MODEL_FILE)
        torch.save(model.state_dict(), model_path)


def train_epoch(model, loss_function, optimizer, data_loader):
    """Train for a single epoch"""
    # set model to training mode
    model.train()

    current_loss = 0.0
    current_acc = 0

    # iterate over the training data
    for i, (inputs, labels) in enumerate(data_loader):

        # zero the parameter gradients
        optimizer.zero_grad()

        with torch.set_grad_enabled(True):
            # forward
            outputs = model(inputs)
            # print(outputs.size(), inputs.size())
            _, predictions = torch.max(outputs, 1)
            loss = loss_function(outputs, labels)

            # backward
            loss.backward()
            optimizer.step()

        # statistics
        current_loss += loss.item() * inputs.size(0)
        current_acc += torch.sum(predictions == labels.data)

    total_loss = current_loss / len(data_loader.dataset)
    total_acc = current_acc.double() / len(data_loader.dataset)

    print('Train Loss: {:.4f}; Accuracy: {:.4f}'.format(total_loss, total_acc))


def test(model, loss_function, data_loader):
    """Test over the whole dataset"""

    model.eval()  # set model in evaluation mode

    current_loss = 0.0
    current_acc = 0

    # iterate over the validation data
    for i, (inputs, labels) in enumerate(data_loader):
        # forward
        with torch.set_grad_enabled(False):
            outputs = model(inputs)
            _, predictions = torch.max(outputs, 1)
            loss = loss_function(outputs, labels)

        # statistics
        current_loss += loss.item() * inputs.size(0)
        current_acc += torch.sum(predictions == labels.data)

    total_loss = current_loss / len(data_loader.dataset)
    total_acc = current_acc.double() / len(data_loader.dataset)

    print('Test Loss: {:.4f}; Accuracy: {:.4f}'
          .format(total_loss, total_acc))


if __name__ == '__main__':
    print('Training...')
    m = Net()
    m.eval()
    train(m)
    print('Training Done!')
    x_ex = create_ex_datasets()

    print('Outputs of Neural Network are as follows:')

    for i, (input, label) in enumerate(x_ex):
        print("Example:",i+1)
        output = m(input)
        print(output.tolist()[0])