Alexnet draft PyTorch

alex_torch.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data as data

import skimage
import skimage.io
import skimage.transform

import numpy as np


def load_traffic_sign_dataset_no_resize():
    import os

    def load_data(directory):
        directories = [d for d in os.listdir(directory)
                       if os.path.isdir(os.path.join(directory, d))]
        labels = []
        images = []
        for d in directories:
            label_directory = os.path.join(directory, d)
            file_names = [os.path.join(label_directory, f)
                          for f in os.listdir(label_directory)
                          if f.endswith(".ppm")]
            for f in file_names:
                images.append(skimage.io.imread(f))
                labels.append(int(d))
        images, labels = np.array(images), np.array(labels)
        images = np.array([skimage.transform.resize(img, (128, 128)) for img in images])
        return images, labels

    X_train, y_train = load_data('../assignments/assignment_9/BelgiumTSC_Training/Training')
    X_test, y_test = load_data('../assignments/assignment_9/BelgiumTSC_Testing/Testing')

    return X_train, X_test, y_train, y_test


class TrafficSignCNN:
    def __init__(self, input_shape=(128, 128, 3)):
        self.input_shape = input_shape
        self.n_classes = 61
        self.model = None

        self.X = None
        self.y = None

        self.build_model()

    # alexnet
    def build_model(self):
        self.model = nn.Sequential(
            # conv layer 1
            nn.Conv2d(
                kernel_size=(11, 11),
                stride=(4, 4),
                in_channels=3,
                # also known as number of filters
                out_channels=96
            ),
            nn.BatchNorm2d(96),
            nn.ReLU(),
            nn.MaxPool2d(
                kernel_size=(2, 2),
                stride=(2, 2)
            ),

            # conv layer 2
            nn.Conv2d(
                kernel_size=(5, 5),
                stride=(1, 1),
                in_channels=96,
                # also known as number of filters
                out_channels=256
            ),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(
                kernel_size=(2, 2),
                stride=(2, 2)
            ),

            # conv layer 3
            nn.Conv2d(
                kernel_size=(3, 3),
                stride=(1, 1),
                in_channels=256,
                # also known as number of filters
                out_channels=384,
                padding=(3, 3)
            ),
            nn.BatchNorm2d(384),
            nn.ReLU(),

            # conv layer 4
            nn.Conv2d(
                kernel_size=(3, 3),
                stride=(1, 1),
                in_channels=384,
                # also known as number of filters
                out_channels=384
            ),
            nn.BatchNorm2d(384),
            nn.ReLU(),

            # conv layer 5
            nn.Conv2d(
                kernel_size=(3, 3),
                stride=(1, 1),
                in_channels=384,
                # also known as number of filters
                out_channels=256
            ),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(
                kernel_size=(2, 2),
                stride=(2, 2)
            ),

            # fully connected layers
            # layer 1
            nn.Flatten(),
            nn.Linear(in_features=4096, out_features=4096),
            nn.BatchNorm2d(4096),
            nn.ReLU(),
            nn.Dropout(p=0.4),

            # layer 2
            nn.Linear(out_features=4096, in_features=1000),
            nn.BatchNorm2d(1000),
            nn.ReLU(),
            nn.Dropout(p=0.4),

            # layer 3
            nn.Linear(out_features=self.n_classes, in_features=1000),
            nn.BatchNorm2d(self.n_classes),
            nn.ReLU(),
            nn.Dropout(p=0.4),

            # output layer
            nn.Linear(in_features=self.n_classes, out_features=self.n_classes),
            nn.BatchNorm2d(self.n_classes),
            nn.Softmax2d()
        )

        print(self.model)

    def forward(self, X):
        t_res = self.model(X.float())
        return t_res

    def fit(self, X, y):
        self.X = torch.from_numpy(X)
        self.y = torch.from_numpy(y)

    def train(self, epochs=5):
        print(f"Training model over {epochs} epochs.")
        loss = nn.CrossEntropyLoss()
        optimizer = optim.Adam(self.model.parameters(), lr=3e-4, weight_decay=0.001)
        for epoch in range(epochs):
            optimizer.zero_grad()
            output = self.forward(self.X)
            l = loss(output, self.y)
            l.backward()
            optimizer.step()
            print(f"epoch = {epoch}\tloss = {l.data}")


print("Loading data")
X_train, X_test, y_train, y_test = load_traffic_sign_dataset_no_resize()
print("Done Loading data.")

X_train = X_train.reshape((4575, 3, 128, 128))


cnn = TrafficSignCNN()
cnn.fit(X_train, y_train)
cnn.train()