!pip install jovian --upgrade -q

!pip install torch===1.7.1+cu110 torchvision===0.8.2+cu110 torchaudio===0.7.2 -f https://download.pytorch.org/whl/torch_stable.html

!pip install seaborn

import jovian
import os
import numpy as np
import torch
import torchvision
import tarfile
import matplotlib
import matplotlib.pyplot as plt
import torchvision.transforms as tt
import torch.nn as nn
import torch.nn.functional as F
import seaborn as sns

from torchvision.datasets import ImageFolder
from torchvision.transforms import ToTensor
from torchvision.datasets.utils import download_url
from torch.utils.data import random_split
from torch.utils.data.dataloader import DataLoader
from torchvision.utils import make_grid

%matplotlib inline


matplotlib.rcParams['figure.facecolor'] = '#ffffff'

project_name='zero-to-gans-course-project2'

dataset= 'faces_dataset'
dataset_url='faces_dataset'
path = 'C:\\Users\\Geralt\\Desktop\\Deeplearning\\zero-to-gans-course-project3\\faces_dataset'

#checking for cuda drivers if available
def get_default_device():
    """Pick GPU if available, else CPU"""
    if torch.cuda.is_available():
        return torch.device('cuda')
    else:
        return torch.device('cpu')
    
def to_device(data, device):
    """Move tensor(s) to chosen device"""
    if isinstance(data, (list,tuple)):
        return [to_device(x, device) for x in data]
    return data.to(device, non_blocking=True)

class DeviceDataLoader():
    """Wrap a dataloader to move data to a device"""
    def __init__(self, dl, device):
        self.dl = dl
        self.device = device
        
    def __iter__(self):
        """Yield a batch of data after moving it to device"""
        for b in self.dl: 
            yield to_device(b, self.device)

    def __len__(self):
        """Number of batches"""
        return len(self.dl)

device = get_default_device()
device

print(os.listdir(dataset))
classes = os.listdir(dataset + "/Training")
print(classes)

data_dir = path #dividing the data into seperate directories

train_dir = data_dir + '/Training'
test_dir = data_dir + '/Testing'

count = []
for folder in classes:
    num_images = len(os.listdir(train_dir+'/'+ folder))
    count.append(num_images)
    print(f'There are {num_images} images in the {folder} category.')

plt.figure(figsize=(12, 6))
sns.barplot(x=classes, y=count)
plt.title('imgs per category', size=16)
plt.ylabel('Imgs', size=14)
plt.xlabel('Categories', size=14)
plt.show;

def denormalize(images, means, stds):
    means = torch.tensor(means).reshape(1, 3, 1, 1)
    stds = torch.tensor(stds).reshape(1, 3, 1, 1)
    return images * stds + means

# Data transforms (normalization & data augmentation)
stats = ((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
 
train_tfms = tt.Compose([
                         tt.Resize(64),
                         tt.RandomCrop(50, padding=4, padding_mode='reflect'), 
                         tt.RandomHorizontalFlip(), 
                         #tt.RandomRotate,
                         #tt.RandomResizedCrop(256, scale=(0.5,0.9), ratio=(1, 1)), 
                         #tt.ColorJitter(brightness=.1, contrast=.8, saturation=0.5, hue=0.12),
                         tt.ToTensor(), 
                         tt.Normalize(*stats,inplace=True)])
valid_tfms = tt.Compose([tt.ToTensor(), tt.Normalize(*stats)])

train_ds = ImageFolder(train_dir, train_tfms)
test_ds = ImageFolder(test_dir, valid_tfms)

def show_example(img, label):
    print(f'Label: {train_ds.classes[label]} ({label})')
    print(f'image.shape: {img.shape}')
    plt.imshow(img.permute(1, 2, 0))

print(show_example(*train_ds[1]))

show_example(*train_ds[4000])

show_example(*train_ds[8927])

show_example(*train_ds[17000])

batch_size= 350

# PyTorch data loaders
train_dl = DataLoader(train_ds, batch_size, shuffle=True, num_workers=6, pin_memory=True)
test_dl = DataLoader(test_ds, batch_size*2, num_workers=6, pin_memory=True)

def accuracy(outputs, labels): #setting up the model architecture
    _, preds = torch.max(outputs, dim=1)
    return torch.tensor(torch.sum(preds == labels).item() / len(preds))

class ImageClassificationBase(nn.Module):
    def training_step(self, batch):
        images, labels = batch 
        out = self(images)                  # Generate predictions
        loss = F.cross_entropy(out, labels) # Calculate loss
        return loss
    
    def validation_step(self, batch):
        images, labels = batch 
        out = self(images)                    # Generate predictions
        loss = F.cross_entropy(out, labels)   # Calculate loss
        acc = accuracy(out, labels)           # Calculate accuracy
        return {'val_loss': loss.detach(), 'val_acc': acc}
        
    def validation_epoch_end(self, outputs):
        batch_losses = [x['val_loss'] for x in outputs]
        epoch_loss = torch.stack(batch_losses).mean()   # Combine losses
        batch_accs = [x['val_acc'] for x in outputs]
        epoch_acc = torch.stack(batch_accs).mean()      # Combine accuracies
        return {'val_loss': epoch_loss.item(), 'val_acc': epoch_acc.item()}
    
    def epoch_end(self, epoch, result):
        print("Epoch [{}], last_lr: {:.5f}, train_loss: {:.4f}, val_loss: {:.4f}, val_acc: {:.4f}".format(
            epoch, result['lrs'][-1], result['train_loss'], result['val_loss'], result['val_acc']))

def conv_block(in_channels, out_channels, pool=False):
    layers = [nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), 
              nn.BatchNorm2d(out_channels), 
              nn.ReLU(inplace=True)]
    if pool: layers.append(nn.MaxPool2d(2))
    return nn.Sequential(*layers)

class ResNet9(ImageClassificationBase): #resnet9 architecture
    def __init__(self, in_channels, num_classes):
        super().__init__()
        # channels: 1x48x48
        
        self.conv1 = conv_block(in_channels, 64) #64x48x48
        self.conv2 = conv_block(64, 128, pool=True)#128x24x24
        self.res1 = nn.Sequential(conv_block(128, 128), #takes 128 imput channels and returns a feature map of the same number 
                                  conv_block(128, 128)) #the output of these two residual blocks will stay the same as the residual block imput so the input can be fed back in
        
        self.conv3 = conv_block(128, 256, pool=True) #256x12x12
        self.conv4 = conv_block(256, 512, pool=True) #512x6x6
        self.res2 = nn.Sequential(conv_block(512, 512),
                                  conv_block(512, 512))#this is the second residual block 
        self.classifier = nn.Sequential(nn.MaxPool2d(4), #performing maxpooling to convert the blocks into the max value 512x1x1
                                        nn.Flatten(), #flatten the 3d map #512
                                        nn.Dropout(0.35), #avoiding overfitting by randomly picking 20% of the feature map elements and setting them to zero. only keeping 80% of the learned data
                                        nn.Linear(512, num_classes)) # takes the 512 outputs and converts them into the outputs (6) - lesson 2
        
    def forward(self, xb):
        out = self.conv1(xb)
        out = self.conv2(out)
        out = self.res1(out) + out #on the residual layer we take the output and add the output from conv2
        out = self.conv3(out)
        out = self.conv4(out)
        out = self.res2(out) + out
        out = self.classifier(out)
        return out

model = ResNet9(3, 10)
model

device = get_default_device()
device

train_dl = DeviceDataLoader(train_dl, device)
test_dl = DeviceDataLoader(test_dl, device)
to_device(model, device);

@torch.no_grad()
def evaluate(model, val_loader):
    model.eval()
    outputs = [model.validation_step(batch) for batch in val_loader]
    return model.validation_epoch_end(outputs)

def get_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']

def fit_one_cycle(epochs, max_lr, model, train_loader, val_loader, 
                  weight_decay=0, grad_clip=None, opt_func=torch.optim.SGD):
    torch.cuda.empty_cache()
    history = []
    
    # Set up cutom optimizer with weight decay
    optimizer = opt_func(model.parameters(), max_lr, weight_decay=weight_decay)
    # Set up one-cycle learning rate scheduler
    sched = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr, epochs=epochs, 
                                                steps_per_epoch=len(train_loader))
    
    for epoch in range(epochs):
        # Training Phase 
        model.train()
        train_losses = []
        lrs = []
        for batch in train_loader:
            loss = model.training_step(batch)
            train_losses.append(loss)
            loss.backward()
            
            # Gradient clipping
            if grad_clip: 
                nn.utils.clip_grad_value_(model.parameters(), grad_clip)
            
            optimizer.step()
            optimizer.zero_grad()
            
            # Record & update learning rate
            lrs.append(get_lr(optimizer))
            sched.step()
        
        # Validation phase
        result = evaluate(model, val_loader)
        result['train_loss'] = torch.stack(train_losses).mean().item()
        result['lrs'] = lrs
        model.epoch_end(epoch, result)
        history.append(result)
    return history

history = [evaluate(model, test_dl)]
history

epochs =20
max_lr = .02
grad_clip = 0.1
weight_decay = 1e-4
opt_func = torch.optim.Adam

%%time
history += fit_one_cycle(epochs, max_lr, model, train_dl, test_dl, 
                             grad_clip=grad_clip, 
                             weight_decay=weight_decay, 
                             opt_func=opt_func)

epochs = 10
max_lr = .008
grad_clip = 0.08
weight_decay = 1e-5
opt_func = torch.optim.Adam

%%time
history += fit_one_cycle(epochs, max_lr, model, train_dl, test_dl, 
                             grad_clip=grad_clip, 
                             weight_decay=weight_decay, 
                             opt_func=opt_func)

epochs = 10
max_lr = .005
grad_clip = 0.07
weight_decay = 1e-7
opt_func = torch.optim.Adam

%%time
history += fit_one_cycle(epochs, max_lr, model, train_dl, test_dl, 
                             grad_clip=grad_clip, 
                             weight_decay=weight_decay, 
                             opt_func=opt_func)

train_time= '43:35'

torch.save(model.state_dict(), 'faces.pth')

jovian.reset()
jovian.log_hyperparams(arch='test 5: resnet9', 
                       epochs=epochs, 
                       lr=max_lr, 
                       scheduler='one-cycle', 
                       weight_decay=weight_decay,
                       grad_clip=grad_clip, 
                       opt=opt_func.__name__)

jovian.log_metrics(val_loss=history[-1]['val_loss'], 
                   val_acc=history[-1]['val_acc'],
                   train_loss=history[-1]['train_loss'],
                   time=train_time)

project_name='zero-to-gans-course-project2'

jovian.commit(project=project_name, filename= "zero-to-gans-course-project3", files=['faces_dataset'])

import random

def denormalize(images, mean, std):
    invTrans = tt.Compose([ tt.Normalize(mean=[ 0., 0., 0. ],
                                                     std = [ 1/std[0], 1/std[1], 1/std[2] ]),
                                tt.Normalize(mean = [-1*mean[0], -1*mean[2], -1*mean[2] ],
                                                     std = [ 1., 1., 1. ]),
                               ])
    return invTrans(images)

def predict_img(img, label):
    xb = to_device(img.unsqueeze(0), device)
    yb = model(xb)
    _, pred = torch.max(yb, dim=1)
    print(f'Label: {train_ds.classes[label]}, Predicted: {train_ds.classes[pred]}')
    img = denormalize(img, *stats)
    plt.imshow(img.permute(1, 2, 0))

predict_img(*test_ds[1])

predict_img(*test_ds[4000])

predict_img(*test_ds[6000])

predict_img(*test_ds[random.randint(6000,7000)])

predict_img(*test_ds[random.randint(0,5000)])

predict_img(*test_ds[random.randint(5000,6000)])

def plot_accuracies(history):
    accuracies = [x['val_acc'] for x in history]
    plt.plot(accuracies, '-x')
    plt.xlabel('epoch')
    plt.ylabel('accuracy')
    plt.title('Accuracy vs. No. of epochs');

def plot_losses(history):
    train_losses = [x.get('train_loss') for x in history]
    val_losses = [x['val_loss'] for x in history]
    plt.plot(train_losses, '-bx')
    plt.plot(val_losses, '-rx')
    plt.xlabel('epoch')
    plt.ylabel('loss')
    plt.legend(['Training', 'Validation'])
    plt.title('Loss vs. No. of epochs');

def plot_lrs(history):
    lrs = np.concatenate([x.get('lrs', []) for x in history])
    plt.plot(lrs)
    plt.xlabel('Batch no.')
    plt.ylabel('Learning rate')
    plt.title('Learning Rate vs. Batch no.');

plot_accuracies(history)

plot_losses(history)

plot_lrs(history)

jovian.commit(project=project_name, files='facial-recognition-dataset', outputs=['faces.pth'])