# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: pytorch
#     name: pytorch
# ---

# + [markdown] id="3cejx5xaJH6G"
# # Cloning <a  href='https://github.com/nikshrimali/shrinkai'>SHRINKai</a>

# + id="39t2TNs7EYCU" tags=[] outputId="ab907a7c-dd11-43db-e96d-a9a66588c928" colab={"base_uri": "https://localhost:8080/", "height": 118}
# Cloning shrinkai
# # !git clone https://github.com/nikshrimali/shrinkai

# + [markdown] id="L_RUchliJB-_"
# # Importing Libraries

# + id="PKulaUm9Hriq" tags=[] outputId="e4b74027-ee37-423d-f6ac-2094c2427657" colab={"base_uri": "https://localhost:8080/", "height": 34}

import torch
import torchvision
import torchvision.transforms as transforms
from torchsummary import summary
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR, ReduceLROnPlateau


from shrinkai.model.model_test import model_testing
from shrinkai.model.model_train import model_training

from shrinkai.data_process.getdata import GetCIFAR10_TrainData
from shrinkai.data_process.misclassified_data import *

import matplotlib.pyplot as plt
# %matplotlib inline


import warnings
warnings.filterwarnings('ignore')

device = "cuda" if torch.cuda.is_available else "cpu"
print(device)

from shrinkai.model.gradcam import gen_gradcam

# + tags=[] id="bihvtVPj9tk6" outputId="4b272917-ee54-4d80-d56b-ab2b60ff30fd" colab={"base_uri": "https://localhost:8080/", "height": 346}
from shrinkai.data_process.albumentation import *

trainset = datasets.CIFAR10('./data', train=True, download=True, transform=cifar_alb_trainData())
testset = datasets.CIFAR10('./data', train=False, download=True, transform=cifar_alb_testdata())

trainloader = torch.utils.data.DataLoader(trainset, batch_size=128,
                                          shuffle=True, num_workers=0)
testloader = torch.utils.data.DataLoader(testset, batch_size=128,
                                         shuffle=False, num_workers=0)

# + [markdown] id="HZFXAW9CI5yt"
# # Importing Model

# + id="pul6T5YZHyjB" tags=[]
from shrinkai.model.resnetmodel8 import ResNet18
model = ResNet18().to(device)
summary(model, input_size=(3,32,32))

# + id="aJ68Gzm-IFIq" tags=[]
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
# scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.05, patience=2, threshold=0.001, threshold_mode='rel', cooldown=0, min_lr=0, eps=1e-08, verbose=True)
scheduler = StepLR(optimizer, step_size=25, gamma=0.1)

train_acc = []
train_losses = []
test_acc = []
test_losses = []
EPOCHS = 20
model_path = 'latest_model.h5'

import os

print(f'Starting Training for {EPOCHS} Epochs')
try:
    os.remove(model_path) # deleting the existing file
except:
    pass

for i in range(EPOCHS):
    print(f'EPOCHS : {i}')
    model_training(model, device, trainloader, optimizer, train_acc, train_losses, l1_loss=False)
    torch.save(model.state_dict(), model_path)
    scheduler.step()
    misclassified = model_testing(model, device, testloader, test_acc, test_losses)

# + [markdown] id="VYRArZkVJw5C"
# # Model Performance

# + id="4J2uJUfl9tli"
plt.plot(train_acc)
plt.plot(test_acc)

# + id="D26lHx9w9tlm"
from model.gradcam.utils import plot_gradcam

class_names = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
target_layers = ["layer1", "layer2", "layer3", "layer4"]
mean = (0.491, 0.482, 0.446)
std = (0.247, 0.243, 0.261)

plot_gradcam(target_layers, device, testloader, model, mean, std, class_names)

StarcoderdataJupyterScriptsDedupFiltered