src/model_obj.py

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

Automatically generated by Colaboratory.

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




import torch
import torch.optim as optim

from src.model_yolov3 import YOLOv3
from tqdm import tqdm
import src.config as config
from src.utils_rh import (
    mean_average_precision,
    cells_to_bboxes,
    get_evaluation_bboxes,
    save_checkpoint,
    load_checkpoint,
    check_class_accuracy,
    get_loaders,
    plot_couple_examples
)
from src.loss import YoloLoss
import warnings
warnings.filterwarnings("ignore")

import torch
from pytorch_lightning import LightningModule, Trainer
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader, random_split
import torchvision
from pytorch_lightning.callbacks import LearningRateMonitor
from pytorch_lightning.callbacks.progress import TQDMProgressBar
from pytorch_lightning.loggers import CSVLogger
from pytorch_lightning.callbacks import ModelCheckpoint
import pandas as pd
from torch.optim.lr_scheduler import OneCycleLR

import seaborn as sn

class Assignment13(LightningModule):
    def __init__(self):
        super().__init__()
        self.save_hyperparameters()
        self.epoch_number = 0
        self.config = config
        self.train_csv_path = self.config.DATASET + "/train.csv"
        self.test_csv_path = self.config.DATASET + "/test.csv"
        self.train_loader, self.test_loader, self.train_eval_loader = get_loaders(
              train_csv_path=self.train_csv_path, test_csv_path=self.test_csv_path)
        self.check_class_accuracy = check_class_accuracy
        self.model = YOLOv3(num_classes=self.config.NUM_CLASSES)
        self.loss_fn = YoloLoss()
        self.check_class_accuracy = check_class_accuracy
        self.get_evaluation_bboxes = get_evaluation_bboxes
        self.scaled_anchors = (torch.tensor(self.config.ANCHORS) * torch.tensor(self.config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2))
        self.losses = []
        self.plot_couple_examples = plot_couple_examples
        self.mean_average_precision = mean_average_precision
        self.EPOCHS = self.config.NUM_EPOCHS * 2 // 5
    def forward(self, x):
        out = self.model(x)
        return out
    def training_step(self, batch, batch_idx):
        x, y = batch
        out = self(x)
        y0, y1, y2 = (y[0],y[1],y[2])
        loss = (
                self.loss_fn(out[0], y0, self.scaled_anchors[0].to(y0))
                + self.loss_fn(out[1], y1, self.scaled_anchors[1].to(y1))
                + self.loss_fn(out[2], y2, self.scaled_anchors[2].to(y2))
            )
        self.losses.append(loss.item())
        mean_loss = sum(self.losses) / len(self.losses)
        self.log("train_loss", mean_loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
        #self.log("train_loss", mean_loss)
        return loss


    def on_train_epoch_start(self):
        self.epoch_number += 1
        self.losses = []
        #self.plot_couple_examples(self.model,self.test_loader,0.6,0.5,self.scaled_anchors)
        if self.epoch_number > 1 and self.epoch_number % 10 == 0:
            self.plot_couple_examples(self.model,self.test_loader,0.6,0.5,self.scaled_anchors)

    def on_train_epoch_end(self):
        print(f"Currently epoch {self.epoch_number}")
        print("On Train Eval loader:")
        print("On Train loader:")
        self.check_class_accuracy(self.model, self.train_loader, threshold=self.config.CONF_THRESHOLD)
        if self.epoch_number == self.EPOCHS:
              #if self.epoch_number > 1 and self.epoch_number % 3 == 0:
            self.check_class_accuracy(self.model, self.test_loader, threshold=self.config.CONF_THRESHOLD)
            pred_boxes, true_boxes = self.get_evaluation_bboxes( self.test_loader,self.model,iou_threshold=self.config.NMS_IOU_THRESH,
                                                                 anchors=self.config.ANCHORS,
                                                                 threshold=self.config.CONF_THRESHOLD,)
            mapval = self.mean_average_precision(
                pred_boxes,
                true_boxes,
                iou_threshold=self.config.MAP_IOU_THRESH,
                box_format="midpoint",
                num_classes=self.config.NUM_CLASSES,
            )
            print(f"MAP: {mapval.item()}")
            self.model.train()
            pass


    def configure_optimizers(self):
        optimizer = optimizer = optim.Adam(
                    model.parameters(), lr=config.LEARNING_RATE, weight_decay=config.WEIGHT_DECAY)
        #EPOCHS = config.NUM_EPOCHS * 2 // 5
        scheduler = OneCycleLR(
        optimizer,
        max_lr=1E-3,
        steps_per_epoch=len(self.train_loader),
        epochs=self.EPOCHS,
        pct_start=5/self.EPOCHS,
        div_factor=100,
        three_phase=False,
        final_div_factor=100,
        anneal_strategy='linear'
        )

        return {"optimizer": optimizer, "lr_scheduler":scheduler}

     ####################
    # DATA RELATED HOOKS
    ####################

    def train_dataloader(self):
          return self.train_loader

    def test_dataloader(self):
          return self.test_loader

#finding maximum learning rate
model = Assignment13()
#trainer = Trainer(precision=16,accelerator='cpu',callbacks=[TQDMProgressBar(refresh_rate=0)])

# Run learning rate finder
#lr_finder = trainer.tuner.lr_find(model,max_lr=2, num_training=200,mode="exponential")

# Inspect results
#fig = lr_finder.plot(); fig.show()
#suggested_lr = lr_finder.suggestion()
#print(suggested_lr)
# Overwrite lr and create new model
#hparams.lr = suggested_lr
#model = MyModelClass(hparams)

class Assignment13(LightningModule):
    def __init__(self):
        super().__init__()
        self.save_hyperparameters()
        self.epoch_number = 0
        self.config = config
        self.train_csv_path = self.config.DATASET + "/train.csv"
        self.test_csv_path = self.config.DATASET + "/test.csv"
        self.train_loader, self.test_loader, self.train_eval_loader = get_loaders(
              train_csv_path=self.train_csv_path, test_csv_path=self.test_csv_path)
        self.check_class_accuracy = check_class_accuracy
        self.model = YOLOv3(num_classes=self.config.NUM_CLASSES)
        self.loss_fn = YoloLoss()
        self.check_class_accuracy = check_class_accuracy
        self.get_evaluation_bboxes = get_evaluation_bboxes
        self.scaled_anchors = (torch.tensor(self.config.ANCHORS) * torch.tensor(self.config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2))
        self.losses = []
        self.plot_couple_examples = plot_couple_examples
        self.mean_average_precision = mean_average_precision
        self.EPOCHS = self.config.NUM_EPOCHS * 2 // 5
    def forward(self, x):
        out = self.model(x)
        return out
    def training_step(self, batch, batch_idx):
        x, y = batch
        out = self(x)
        y0, y1, y2 = (y[0],y[1],y[2])
        loss = (
                self.loss_fn(out[0], y0, self.scaled_anchors[0].to(y0))
                + self.loss_fn(out[1], y1, self.scaled_anchors[1].to(y1))
                + self.loss_fn(out[2], y2, self.scaled_anchors[2].to(y2))
            )
        self.losses.append(loss.item())
        mean_loss = sum(self.losses) / len(self.losses)
        self.log("train_loss", mean_loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
        #self.log("train_loss", mean_loss)
        return loss


    def on_train_epoch_start(self):
        self.epoch_number += 1
        self.losses = []
        #self.plot_couple_examples(self.model,self.test_loader,0.6,0.5,self.scaled_anchors)
        if self.epoch_number > 1 and self.epoch_number % 10 == 0:
            self.plot_couple_examples(self.model,self.test_loader,0.6,0.5,self.scaled_anchors)

    def on_train_epoch_end(self):
        print(f"Currently epoch {self.epoch_number}")
        print("On Train Eval loader:")
        print("On Train loader:")
        self.check_class_accuracy(self.model, self.train_loader, threshold=self.config.CONF_THRESHOLD)
        if self.epoch_number == self.EPOCHS:
              #if self.epoch_number > 1 and self.epoch_number % 3 == 0:
            self.check_class_accuracy(self.model, self.test_loader, threshold=self.config.CONF_THRESHOLD)
            pred_boxes, true_boxes = self.get_evaluation_bboxes( self.test_loader,self.model,iou_threshold=self.config.NMS_IOU_THRESH,
                                                                 anchors=self.config.ANCHORS,
                                                                 threshold=self.config.CONF_THRESHOLD,)
            mapval = self.mean_average_precision(
                pred_boxes,
                true_boxes,
                iou_threshold=self.config.MAP_IOU_THRESH,
                box_format="midpoint",
                num_classes=self.config.NUM_CLASSES,
            )
            print(f"MAP: {mapval.item()}")
            self.model.train()
            pass


    def configure_optimizers(self):
        optimizer = optimizer = optim.Adam(
                    model.parameters(), lr=config.LEARNING_RATE, weight_decay=config.WEIGHT_DECAY)
        #EPOCHS = config.NUM_EPOCHS * 2 // 5
        scheduler = OneCycleLR(
        optimizer,
        max_lr=8E-4,
        steps_per_epoch=len(self.train_loader),
        epochs=self.EPOCHS,
        pct_start=5/self.EPOCHS,
        div_factor=100,
        three_phase=False,
        final_div_factor=100,
        anneal_strategy='linear'
        )

        return {"optimizer": optimizer, "lr_scheduler":scheduler}

     ####################
    # DATA RELATED HOOKS
    ####################

    def train_dataloader(self):
          return self.train_loader

    def test_dataloader(self):
          return self.test_loader

model = Assignment13()
checkpoint_callback = ModelCheckpoint(
    monitor='train_loss',  # Metric to monitor for saving the best model
    mode='min', # 'min' to save the model with the lowest value of the monitored metric
    dirpath='/storage/',
    filename='assignment13_final{epoch:02d}-train_loss_min_A400{train_loss:.2f}',
    save_top_k=1         # Save only the best model
)

#trainer = Trainer(
    #max_epochs=config.NUM_EPOCHS * 2 // 5,
    #accelerator="cpu",
    #precision=16,  # limiting got iPython runs
    #logger=CSVLogger(save_dir="logs/"),
    #callbacks=[LearningRateMonitor(logging_interval="step"),TQDMProgressBar(refresh_rate=0),checkpoint_callback],
#)


#trainer.fit(model)

#metrics = pd.read_csv(f"{trainer.logger.log_dir}/metrics.csv")
#del metrics["step"]
#metrics.set_index("epoch", inplace=True)
#display(metrics.dropna(axis=1, how="all").head())
#sn.relplot(data=metrics, kind="line")