# -*- coding: utf-8 -*-
"""proj3_federico.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1tC-fSk11g3KiUbQj30K0xTgH9riDKkdp
"""

import torch 
import torchvision
import torch.nn as nn
import numpy as np
import torchvision.transforms as transforms
import pickle
import time

# Device configuration
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Hyper parameters
num_epochs = 300
num_classes = 10
batch_size = 64
#batch_size = 256 #(change 5A)
learning_rate = 0.001
#learning_rate = 0.1 #(change 5B)
leaky_relu_slope = 0.2
#droupout_rate = 0.3

normalize = transforms.Normalize(mean=[.5, .5, .5],
                                 std=[1 ,1, 1])
transform_ = transforms.Compose([
    transforms.ToTensor(),
    normalize,
])

# Download and construct CIFAR-10 dataset.
train_dataset = torchvision.datasets.CIFAR10(root='../../data/',
                                             train=True, 
                                             transform=transform_,
                                             download=True)

image, label = train_dataset[44]
print (image.size(0))
# print (label)
# # print(train_dataset.)

# Data loader (this provides queues and threads in a very simple way).
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,   
                                           batch_size=batch_size, 
                                           shuffle=True,
                                          )

# data_iter = iter(train_loader)
# images, labels = data_iter.next()

test_dataset = torchvision.datasets.CIFAR10(root='../../data/',
                                          train=False, 
                                          transform=transform_,
                                          )

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size, 
                                          shuffle=False)

# Convolutional neural network (two convolutional layers)
class ConvNet(nn.Module):
    def __init__(self, num_classes=10):
        super(ConvNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=5, stride=1, padding=0),
            nn.LeakyReLU(leaky_relu_slope, inplace=True), #leakyrelu instead of Rely (change C)
            nn.BatchNorm2d(64), #batch normalization (change E)
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=0),
            nn.LeakyReLU(leaky_relu_slope, inplace=True), #leakyrelu instead of Rely (change C)
            nn.BatchNorm2d(128), #batch normalization (change E)
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer3 = nn.Sequential(
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=0),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.fc1 = nn.Sequential(
            nn.Linear(2*2*256, 512),
            nn.LeakyReLU(leaky_relu_slope, inplace=True)) #leakyrelu instead of Rely (change C)
        #self.dropout = nn.Dropout(dropout_rate) #add dropout with rate=0.3 (change D)
        self.fc2 = nn.Linear(512, num_classes)
        self.sm = nn.Softmax(dim=-1)
        
    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc1(out)
        #out = self.dropout(out) #(change D)
        out = self.fc2(out)
        out = self.sm(out)
        return out

model = ConvNet(num_classes).to(device)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate) #if needed momentum=0.9

# Train the model
losses = []
total_step = len(train_loader) #number of batches

#start training time counting
torch.cuda.synchronize()
since = int(round(time.time()*1000))

for epoch in range(num_epochs):
    running_loss=0.0
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)
        
        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * images.size(0)
        
        if (i+1) % 100 == 0:
            print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 
                   .format(epoch+1, num_epochs, i+1, total_step, loss.item())) 
    epoch_loss = running_loss / total_step #loss per epoch
    losses.append(epoch_loss)

#stop time counting
torch.cuda.synchronize()
time_elapsed = int(round(time.time()*1000)) - since
print ('training time elapsed {}ms'.format(time_elapsed)) #print time

pickle_out = open("dict.pickle","wb")
pickle.dump(losses, pickle_out)
pickle_out.close()

# Test the model
model.eval()  # eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance)
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    print('Test Accuracy of the model on the 10000 test images: {} %'.format(100 * correct / total))

# Save the model checkpoint
torch.save(model.state_dict(), 'model.ckpt')