!pip3 install 'torch==1.3.1'
!pip3 install 'torchvision==0.4.2'
!pip3 install 'Pillow-SIMD'
!pip3 install 'tqdm'



import os
import logging

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Subset, DataLoader, Dataset
from torch.backends import cudnn

import math

import torchvision
from torchvision import transforms
from torchvision.models import alexnet

from PIL import Image
from tqdm import tqdm
from torchvision.datasets import VisionDataset

import warnings
warnings.filterwarnings("ignore")

import matplotlib.pyplot as plt


DEVICE = 'cuda' # 'cuda' or 'cpu'

NUM_CLASSES = 101 # 101 + 1: There is am extra Background class that should be removed

BATCH_SIZE = 256     # Higher batch sizes allows for larger learning rates. An empirical heuristic suggests that, when changing
                     # the batch size, learning rate should change by the same factor to have comparable results

LR = 0.01 # [0.01,0.02,0.03,0.04,0.05,0.06,0.07,0.08,0.09,0.1] # The initial Learning Rate

from random import seed
from random import random


seed(42) # because the reply of the foundamental question is 42



list_accuracy_param = []

for i in range(0,1):

  MOMENTUM = 0.9       # Hyperparameter for SGD, keep this at 0.9 when using SGD
  WEIGHT_DECAY = 5e-5  # Regularization, you can keep this at the default

  NUM_EPOCHS_O = 15    # Total number of training epochs (iterations over dataset)
  STEP_SIZE_O = 12      # How many epochs before decreasing learning rate (if using a step-down policy)
  GAMMA = 0.1          # Multiplicative factor for learning rate step-down

  LOG_FREQUENCY = 10

  STEP_SIZE = STEP_SIZE_O
  NUM_EPOCHS = NUM_EPOCHS_O

  print(("running with step_size: %f, and num epochs: %d") % (STEP_SIZE,NUM_EPOCHS))

  # Define transforms for training phase
  train_transform = transforms.Compose([transforms.Resize(256),      # Resizes short size of the PIL image to 256
                                        transforms.CenterCrop(224),  # Crops a central square patch of the image
                                                                    # 224 because torchvision's AlexNet needs a 224x224 input!
                                                                    # Remember this when applying different transformations, otherwise you get an error
                                        transforms.ToTensor(), # Turn PIL Image to torch.Tensor
                                        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) # Normalizes tensor with mean and standard deviation
  ])
  # Define transforms for the evaluation phase
  eval_transform = transforms.Compose([transforms.Resize(256),
                                        transforms.CenterCrop(224),
                                        transforms.ToTensor(),
                                        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
  ])

  !rm -rf 'Homework2-Caltech101'

  # Clone github repository with data
  if not os.path.isdir('./Homework2-Caltech101'):
    !git clone https://github.com/eliax1996/Homework2-Caltech101

  !mv 'Homework2-Caltech101' 'Homework2_Caltech101'

  from Homework2_Caltech101 import caltech_dataset as cal_ds

  DATA_DIR = 'Homework2-Caltech101/101_ObjectCategories'

  # Prepare Pytorch train/test Datasets
  train_dataset = cal_ds.Caltech(DATA_DIR, split = 'train', transform=train_transform)
  test_dataset = cal_ds.Caltech(DATA_DIR, split = 'test', transform=eval_transform)


  train_indexes_ = [idx for idx in range(len(train_dataset)) if idx % 5]
  test_indexes = [idx for idx in range(len(test_dataset)) if not idx % 5]


  j = 0
  val_indexes = []
  train_indexes = []
  for ind in train_indexes_:
    if j % 2: val_indexes.append(ind)
    else: train_indexes.append(ind)
    j = j+1


  train_dataset = Subset(train_dataset, train_indexes)
  test_dataset = Subset(test_dataset, test_indexes)
  val_dataset = Subset(train_dataset, val_indexes)

  # Check dataset sizes
  #print('Train Dataset: {}'.format(len(train_dataset)))
  #print('Validation Dataset: {}'.format(len(val_dataset)))
  #print('Test Dataset: {}'.format(len(test_dataset)))


  # Dataloaders iterate over pytorch datasets and transparently provide useful functions (e.g. parallelization and shuffling)
  train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, drop_last=True)
  val_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4)
  test_dataloader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4)


  net = alexnet(pretrained=True) # Loading AlexNet model
  net.classifier[6] = nn.Linear(4096, NUM_CLASSES)

  #sys.exit(0)


  criterion = nn.CrossEntropyLoss()
  parameters_to_optimize = net.parameters()

  print(parameters_to_optimize)

  i = 0
  new_p = []
  for param in parameters_to_optimize:
    if i in (0,3,6,8,10): #(1,4,6)
      param.requires_grad = False
    new_p.append(param)
    i = i+1

  # fully connected
  #net.classifier[1].requires_grad = False
  #net.classifier[4].requires_grad = False
  #net.classifier[6].requires_grad = False

  # convolutional
  #net.features[0].requires_grad = False
  #net.features[3].requires_grad = False
  #net.features[6].requires_grad = False
  #net.features[8].requires_grad = False
  #net.features[10].requires_grad = False


  optimizer = optim.SGD(new_p, lr=LR, momentum=MOMENTUM, weight_decay=WEIGHT_DECAY)
  scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=STEP_SIZE, gamma=GAMMA)



  !rm -rf "weights_only.pth"

  # By default, everything is loaded to cpu
  net = net.to(DEVICE) # this will bring the network to GPU if DEVICE is cuda

  cudnn.benchmark # Calling this optimizes runtime

  current_step = 0

  best_acc = 0
  loss_list = []
  acc_list = []
  # Start iterating over the epochs
  for epoch in range(NUM_EPOCHS):
    print('Starting epoch {}/{}, LR = {}'.format(epoch+1, NUM_EPOCHS, scheduler.get_lr()))

    # Iterate over the dataset
    for images, labels in train_dataloader:
      # Bring data over the device of choice
      images = images.to(DEVICE)
      labels = labels.to(DEVICE)

      net.train() # Sets module in training mode

      # PyTorch, by default, accumulates gradients after each backward pass
      # We need to manually set the gradients to zero before starting a new iteration
      optimizer.zero_grad() # Zero-ing the gradients

      # Forward pass to the network
      outputs = net(images)

      # Compute loss based on output and ground truth
      loss = criterion(outputs, labels)

      # Log loss
      if current_step % LOG_FREQUENCY == 0:
        print('Step {}, Loss {}'.format(current_step, loss.item()))
        loss_list.append(loss)

      if math.isnan(loss):
        break

      # Compute gradients for each layer and update weights
      loss.backward()  # backward pass: computes gradients
      optimizer.step() # update weights based on accumulated gradients

      current_step += 1

    # Step the scheduler
    scheduler.step()

    running_corrects=0
    for img,lab in val_dataloader:
      lab = lab.to(DEVICE)
      img = img.to(DEVICE)
      outputs = net(img)
      _, preds = torch.max(outputs.data, 1)
      running_corrects += torch.sum(preds == lab.data).data.item()
      accuracy = running_corrects / float(len(val_dataset))
    acc_list.append(accuracy)

    if accuracy > best_acc:
      print("best accuracy on Validation now is: %f", accuracy)
      best_acc = accuracy
      torch.save(net.state_dict(), 'weights_only.pth')



  plt.plot(loss_list)
  plt.title('loss')
  plt.grid()
  plt.show()
  plt.plot(acc_list)
  plt.title('accuracy')
  plt.grid()
  plt.show()
  print()
  print("the best accuracy on validation set is: %f" % best_acc)
  print()
  print()


  list_accuracy_param.append([[NUM_EPOCHS,STEP_SIZE],best_acc])



  net.load_state_dict(torch.load('weights_only.pth'))

  net.to(DEVICE) # this will bring the network to GPU if DEVICE is cuda
  net.train(False) # Set Network to evaluation mode

  running_corrects = 0
  for images, labels in tqdm(test_dataloader):
    images = images.to(DEVICE)
    labels = labels.to(DEVICE)

    # Forward Pass
    outputs = net(images)

    # Get predictions
    _, preds = torch.max(outputs.data, 1)

    # Update Corrects
    running_corrects += torch.sum(preds == labels.data).data.item()

  # Calculate Accuracy
  accuracy = running_corrects / float(len(test_dataset))

  print('Test Accuracy: {}'.format(accuracy))

print(list_accuracy_param)