hacktag_sideview_modeltrain.py

# -*- coding: utf-8 -*-
"""hack_tag sideview

Automatically generated by Colaboratory.

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

# Commented out IPython magic to ensure Python compatibility.
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import os 
import torchvision
import torch.optim as optim
from torchvision.datasets import ImageFolder
import torchvision.transforms as T
from torch.utils.data import random_split
from torch.utils.data import Dataset,DataLoader
from torchvision.utils import make_grid
import torch.nn.functional as F
from torchvision import models
from tqdm.notebook import tqdm
# %matplotlib inline

!pip install kaggle
!mkdir ~/.kaggle
!cp kaggle.json ~/.kaggle/
!chmod 600 ~/.kaggle/kaggle.json

! kaggle competitions download -c state-farm-distracted-driver-detection

from google.colab import drive
drive.mount('/content/gdrive')

!unzip /content/state-farm-distracted-driver-detection.zip

# Train path
train_path = "/content/imgs/train"
# Test path
test_path = "/content/imgs/test"

train_length = 0
for clas in os.listdir(train_path):
    train_length += len(os.listdir(os.path.join(train_path, clas)))

# Data Transforms and Augmentation
train_transforms = T.Compose([ T.Resize((64, 64)),
                               T.RandomAdjustSharpness(2),
                               T.RandomRotation((-15, 15)),
                               T.ColorJitter(brightness=.5, hue=.3),
                               T.ToTensor(),
                             ])

# Loading Data using ImageFolder
train_ds = ImageFolder(train_path, train_transforms)
classes = train_ds.classes

# Splitting into train-val set
val_pct = .1
val_size = int(val_pct * len(train_ds))
train_ds ,valid_ds = random_split(train_ds, [len(train_ds)-val_size, val_size])

# Data Loader
batch_size = 64
train_dl = DataLoader(train_ds, batch_size, shuffle = True, num_workers=2, pin_memory=True)
valid_dl = DataLoader(valid_ds, batch_size, num_workers=2, pin_memory=True)



class ResNet34(nn.Module):
    def __init__(self,num_classes,pretrained=True):
        super().__init__()
        # Use a pretrained model
        self.network = models.resnet34(pretrained=pretrained)
        # Replace last layer
        self.network.fc = nn.Linear(self.network.fc.in_features, num_classes)

    def forward(self,x):
        return self.network(x)

# Device configuration
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

# Model configuration
pretrained_model = ResNet34(num_classes=10)

pretrained_model.to(device)

# Setting parameters
epochs = 8
lr = 1e-4
opt_func = torch.optim.Adam

def accuracy(outputs,labels):
  _,preds = torch.max(outputs,dim=1)
  return torch.tensor(torch.sum(preds == labels).item() / len(preds))
        
@torch.no_grad()
def evaluate(model, val_loader):
    model.eval()
    outputs = []
    for batch in val_loader:
        images,labels = batch
        images = images.to(device)
        print(images.shape)
        labels = labels.to(device)
        out = model(images)                     
        loss = F.cross_entropy(out, labels)     
        acc = accuracy(out, labels)             
        outputs.append({"val_loss":loss.detach(), "val_acc":acc})
        
    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()}

"""
    Training step
"""
#Set up custom optimizer with weight decay
optimizer = opt_func(pretrained_model.parameters(), lr)

# Saving results for the evaluation of the model and comperison with other model.
history = []

for epoch in range(epochs):
    # Training Phase
    pretrained_model.train()
    train_losses = []
    lrs = []
    for batch in tqdm(train_dl):
        images, labels = batch
        images = images.to(device)
        labels = labels.to(device)
        out = pretrained_model(images)                    
        loss = F.cross_entropy(out, labels)
        train_losses.append(loss)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# Validation Phase
    result = evaluate(pretrained_model,valid_dl)
    result["train_loss"] = torch.stack(train_losses).mean().item()
    # result2["lrs"] = lrs  # I won't add learning rate changes to the history
    print("Epoch [{}/{}], train_loss : {:.4f}, val_loss : {:.4f}], val_acc : {:.4f}".format(epoch, epochs,
                                                                                         result["train_loss"],
                                                                                         result["val_loss"],
                                                                                         result["val_acc"]))
    history.append(result)

torch.save(pretrained_model.state_dict(), "/content/gdrive/MyDrive/Kaam/Hacktag_side_resnet34.pth")
torch.save(model, '/content/gdrive/MyDrive/Kaam/model.pth')

result = evaluate(pretrained_model,valid_dl)