1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

from __future__ import print_function

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from linformer import Linformer
from sklearn.model_selection import train_test_split
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets, transforms
from tqdm.notebook import tqdm
import torchvision
from vit_pytorch.efficient import ViT

1
2
3
4
5
6
7
8

print(f"Is CUDA supported by this system? {torch.cuda.is_available()}")
print(f"CUDA version: {torch.version.cuda}")

# Storing ID of current CUDA device
cuda_id = torch.cuda.current_device()
print(f"ID of current CUDA device: {torch.cuda.current_device()}")

print(f"Name of current CUDA device: {torch.cuda.get_device_name(cuda_id)}")

1
2
3
4
5
6
7
8

batchSize = 128

# Orginial data is list of tuples (PIL Image, class label)
# train_split  = torchvision.datasets.CIFAR100('./cifar-100', train=True,download=True,  transform = transforms.Compose([transforms.ToTensor()]))
# test_split  = torchvision.datasets.CIFAR100('./cifar-100', train=False,download=True,  transform = transforms.Compose([transforms.ToTensor()]))

train_split  = torchvision.datasets.CIFAR100('./cifar-100', train=True,download=True)
test_split  = torchvision.datasets.CIFAR100('./cifar-100', train=False,download=True)

1

train_split[0]

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102

textLabel = [
'apple', # id 0
'aquarium_fish',
'baby',
'bear',
'beaver',
'bed',
'bee',
'beetle',
'bicycle',
'bottle',
'bowl',
'boy',
'bridge',
'bus',
'butterfly',
'camel',
'can',
'castle',
'caterpillar',
'cattle',
'chair',
'chimpanzee',
'clock',
'cloud',
'cockroach',
'couch',
'crab',
'crocodile',
'cup',
'dinosaur',
'dolphin',
'elephant',
'flatfish',
'forest',
'fox',
'girl',
'hamster',
'house',
'kangaroo',
'computer_keyboard',
'lamp',
'lawn_mower',
'leopard',
'lion',
'lizard',
'lobster',
'man',
'maple_tree',
'motorcycle',
'mountain',
'mouse',
'mushroom',
'oak_tree',
'orange',
'orchid',
'otter',
'palm_tree',
'pear',
'pickup_truck',
'pine_tree',
'plain',
'plate',
'poppy',
'porcupine',
'possum',
'rabbit',
'raccoon',
'ray',
'road',
'rocket',
'rose',
'sea',
'seal',
'shark',
'shrew',
'skunk',
'skyscraper',
'snail',
'snake',
'spider',
'squirrel',
'streetcar',
'sunflower',
'sweet_pepper',
'table',
'tank',
'telephone',
'television',
'tiger',
'tractor',
'train',
'trout',
'tulip',
'turtle',
'wardrobe',
'whale',
'willow_tree',
'wolf',
'woman',
'worm',
]

1
2
3
4
5
6
7

random_idx = np.random.randint(1, len(train_split), size=9)
fig, axes = plt.subplots(3, 3, figsize=(16, 12))

for idx, ax in enumerate(axes.ravel()):
    randIndex = random_idx[idx]
    ax.set_title('The label is: ' + textLabel[train_split[randIndex][1]])
    ax.imshow(train_split[randIndex][0])

1
2

print(train_split)
print(test_split)

1

labels = [train_split[i][1] for i in range(len(train_split))]

1

train_list, valid_list = train_test_split(train_split, test_size=0.2, shuffle=True, stratify=labels) #, stratify=[i[1] for i in train_split]

1
2
3
4
5

import plotly.express as px
x = [train_list[i][1] for i in range(len(train_list))]
fig = px.histogram(x)
fig.update_layout(title="Train list",bargap=0.2)
fig.show()

1
2
3
4

x = [valid_list[i][1] for i in range(len(valid_list))]
fig = px.histogram(title="Valid list", y=x)
fig.update_layout(bargap=0.2)
fig.show()

1
2
3
4

x = [test_split[i][1] for i in range(len(test_split))]
fig = px.histogram(x)
fig.update_layout(title="Test list",bargap=0.2)
fig.show()

1
2
3

print(f"Train Data: {len(train_list)}")
print(f"Validation Data: {len(valid_list)}")
print(f"Test Data: {len(test_split)}")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

from torchvision.transforms.autoaugment import AutoAugmentPolicy

all_transforms = transforms.Compose([
                    transforms.RandomHorizontalFlip(),
                    transforms.RandomVerticalFlip(),
                    transforms.RandomRotation(degrees=(0, 180)),
                    transforms.ColorJitter(brightness=0, contrast=0, saturation=0, hue=0),
                    transforms.AutoAugment(AutoAugmentPolicy.CIFAR10),
                    transforms.ToTensor(),
                    # transforms.RandomErasing(),
                ])

val_transforms = transforms.Compose(
    [
        transforms.ToTensor(),
    ]
)

test_transforms = transforms.Compose(
    [
        transforms.ToTensor(),
    ]
)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

class CIFAR100Dataset(Dataset):
    def __init__(self, rawData, transform=None):
        self.rawData = rawData
        self.transform = transform

    def __len__(self):
        self.dataSize = len(self.rawData)
        return self.dataSize

    def __getitem__(self, idx):
        rawData = self.rawData[idx]
        img = rawData[0]
        img_transformed = self.transform(img)
        label = rawData[1]

        return img_transformed, label

1
2
3

train_list_transformed = CIFAR100Dataset(train_list, transform=all_transforms)
valid_list_transformed= CIFAR100Dataset(valid_list, transform=val_transforms)
test_split_transformed= CIFAR100Dataset(test_split, transform=test_transforms)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

random_idx = np.random.randint(1, len(train_list_transformed), size=9)
fig, axes = plt.subplots(3, 3, figsize=(16, 12))
fig.suptitle("Transformed Images", fontsize=14)

for idx, ax in enumerate(axes.ravel()):
    randIndex = random_idx[idx]
    ax.set_title('The label is: ' + textLabel[train_list_transformed[randIndex][1]])
    ax.imshow(transforms.ToPILImage()(train_list_transformed[randIndex][0]))


PATCH_SIZE = 8
PATCH_NUM = int(32 / PATCH_SIZE)
patches = train_list_transformed[random_idx[0]][0].unfold(1, PATCH_SIZE, PATCH_SIZE).unfold(2, PATCH_SIZE, PATCH_SIZE)

fig, ax = plt.subplots(PATCH_NUM, PATCH_NUM)
fig.suptitle("Patched Transformed Images", fontsize=14)
for i in range(PATCH_NUM):
    for j in range(PATCH_NUM):
        sub_img = patches[:, i, j]
        ax[i][j].imshow(torchvision.transforms.functional.to_pil_image(sub_img))
        ax[i][j].axis('off')


patches = patches.reshape(3, -1, PATCH_SIZE, PATCH_SIZE)
patches.transpose_(0, 1)
fig, ax = plt.subplots(1, PATCH_NUM*PATCH_NUM, figsize=(20, 20))
for i in range(PATCH_NUM**2):
    ax[i].imshow(torchvision.transforms.functional.to_pil_image(patches[i]))
    ax[i].axis('off')


fig, ax = plt.subplots(1, 4)
for i in range(4):
    ax[i].imshow(torchvision.transforms.functional.to_pil_image(patches[i]))
    ax[i].axis('off')

1
2
3

train_loader = torch.utils.data.DataLoader(train_list_transformed, batch_size=batchSize, shuffle=True)
valid_loader = torch.utils.data.DataLoader(valid_list_transformed, batch_size=batchSize, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_split_transformed, batch_size=batchSize, shuffle=True)

1

print(len(train_list), len(train_loader))

1

print(len(valid_list), len(valid_loader))

1

train_list_transformed[0]

1
2
3
4
5

train_list_transformed[0][0].shape
ch, seqDim, _ = train_list_transformed[0][0].shape
print(ch, seqDim)
print(train_list_transformed[0][0].shape)
print(len(train_list_transformed))

1
2
3
4
5
6

dim: the dimension of each head in multi-head attention
k: the k that the key/values are projected to along the sequence dimension
heads: number of heads
dropout: the dropout rate for the linear layers
depth: number of transformer block
seq_len: the length of the sequence (number of pixels + class label)

1
2
3
4
5
6
7
8

efficient_transformer = Linformer(
    dim=256,
    seq_len=64+1,  # 8x8 patches + 1 cls-token
    depth=4,
    heads=8,
    k=64,
    dropout = 0.1
)

1

device = 'cuda'

1
2
3
4
5

dim: Last dimension of output tensor after linear transformation
patch_size: Number of patches
image_size: dimension of the input image
num_classes: classes to classify
channels: color channels

1
2
3
4
5
6
7
8

model = ViT(
    dim=256,
    image_size=32,      # 32 pixel by 32 pixel image
    patch_size=4,       # Total 4 patch 8x8 each
    num_classes=100,
    transformer=efficient_transformer,
    channels=3
).to(device)

1
2
3
4
5
6
7

# loss function
criterion = nn.CrossEntropyLoss()
# optimizer
lr = 5e-3
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.99) # weight_decay=0.01
# optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=5e-5) # weight_decay=0.01
scheduler = StepLR(optimizer, step_size=10, gamma=0.8)

1
2

records = []        # A variable to record data for each epoch so we can save the data in a csv later
# model = torch.load('./argumentModel', map_location=device)

1
2

# Waking from suspend cause the nvidia driver to fail sometimes, this command remove and add nvidia_uvm module to solve this problem
# !sudo modprobe -r nvidia_uvm && sudo modprobe nvidia_uvm

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46

scaler = torch.cuda.amp.GradScaler(enabled=True)
model.train()

for epoch in range(200):
    epoch_loss = 0
    epoch_accuracy = 0
    for data, label in tqdm(train_loader):
        data = data.to(device)
        label = label.to(device)

        with torch.autocast(device_type='cuda', dtype=torch.float16):
            output = model(data)
            assert output.dtype is torch.float16
            loss = criterion(output, label)
            assert loss.dtype is torch.float32

        scaler.scale(loss).backward()
        # loss.backward()

        scaler.step(optimizer)
        # scheduler.step()
        # optimizer.step()
        scaler.update()
        optimizer.zero_grad()

        acc = (output.argmax(dim=1) == label).float().mean()
        epoch_accuracy += acc / len(train_loader)
        epoch_loss += loss / len(train_loader)

    with torch.no_grad():
        epoch_val_accuracy = 0
        epoch_val_loss = 0
        for data, label in valid_loader:
            data = data.to(device)
            label = label.to(device)

            val_output = model(data)
            val_loss = criterion(val_output, label)

            acc = (val_output.argmax(dim=1) == label).float().mean()
            epoch_val_accuracy += acc / len(valid_loader)
            epoch_val_loss += val_loss / len(valid_loader)
    records.append([int(epoch+1), epoch_loss.detach().cpu().numpy(), epoch_accuracy.detach().cpu().numpy(), epoch_val_loss.detach().cpu().numpy(), epoch_val_accuracy.detach().cpu().numpy()])
    print(
        f"Epoch : {epoch+1} - loss : {epoch_loss:.4f} - acc: {epoch_accuracy:.4f} - val_loss : {epoch_val_loss:.4f} - val_acc: {epoch_val_accuracy:.4f}\n"
    )

1
2

pytorch_total_params = sum(p.numel() for p in model.parameters())
"Total Model parameters: " + str(pytorch_total_params)

1
2
3
4
5
6
7
8
9

# Remove epoch use index instead
def saveModel(modelName):
    torch.save(model, './' + modelName)
    df = pd.DataFrame(np.array(records), columns=['epoch', 'epoch_loss', 'epoch_accuracy', 'epoch_val_loss', 'epoch_val_accuracy'])
    with open('./' + modelName + '.csv', 'a') as file:
        df.to_csv('./' + modelName + '.csv', mode='a', index=False)
        file.close()
    return df
saveModel('argumentModel3m')

1
2

def getModelCSV(modelName):
    return pd.read_csv('./' + modelName + '.csv')

1
2
3
4
5
6

import hvplot.pandas
def getModelPlot(modelName):
    return getModelCSV(modelName).hvplot(title=f'{modelName}', xlabel='epoch', ylabel='%', use_index=True,
                                         y=['epoch_loss', 'epoch_accuracy', 'epoch_val_loss', 'epoch_val_accuracy'], kind='line')

getModelPlot('argumentModel3m')

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

random_idx = np.random.randint(1, len(valid_list_transformed), size=9)
fig, axes = plt.subplots(3, 3, figsize=(16, 12))
model.eval()

for idx, ax in enumerate(axes.ravel()):
    randIndex = random_idx[idx]
    # input tensor must be [batch size, channels, h, w]
    predictLabel = model(valid_list_transformed[randIndex][0].unsqueeze(0).to(device)).argmax(dim=1)
    trueLabel = valid_list_transformed[randIndex][1]

    ax.set_title('Prediction: ' + textLabel[predictLabel] + '\n'
                 + 'True Label: ' + textLabel[trueLabel])
    ax.imshow(transforms.ToPILImage()(valid_list_transformed[randIndex][0]))