Upload lightning_utils.py

lightning_utils.py

@@ -0,0 +1,171 @@
+
+from torch.utils.data import Dataset, DataLoader
+from loss import YoloLoss
+import config
+import torch
+from dataset import YOLODataset
+from torch.optim.lr_scheduler import OneCycleLR
+import random
+from model import YOLOv3
+import lightning.pytorch as pl
+
+def criterion(out, y, anchors):
+  loss_fn = YoloLoss()
+  loss = (
+          loss_fn(out[0], y[0], anchors[0])
+          + loss_fn(out[1], y[1], anchors[1])
+          + loss_fn(out[2], y[2], anchors[2]))
+  return loss
+
+
+def get_loader(train_dataset, test_dataset):
+  train_loader = DataLoader(
+          dataset=train_dataset,
+          batch_size=config.BATCH_SIZE,
+          num_workers=config.NUM_WORKERS,
+          pin_memory=config.PIN_MEMORY,
+          shuffle=True,
+          drop_last=False,
+  )
+
+  test_loader = DataLoader(
+      dataset=test_dataset,
+      batch_size=config.BATCH_SIZE,
+      num_workers=config.NUM_WORKERS,
+      pin_memory=config.PIN_MEMORY,
+      shuffle=False,
+      drop_last=False,
+      
+  )
+
+  return(train_loader, test_loader)
+
+
+def accuracy_fn(y, out, threshold, 
+                correct_class, correct_obj, 
+                correct_noobj, tot_class_preds, 
+                tot_obj, tot_noobj):
+
+  for i in range(3):
+      
+      obj = y[i][..., 0] == 1 # in paper this is Iobj_i
+      noobj = y[i][..., 0] == 0  # in paper this is Iobj_i
+
+      correct_class += torch.sum(
+          torch.argmax(out[i][..., 5:][obj], dim=-1) == y[i][..., 5][obj]
+      )
+      tot_class_preds += torch.sum(obj)
+
+      obj_preds = torch.sigmoid(out[i][..., 0]) > threshold
+      correct_obj += torch.sum(obj_preds[obj] == y[i][..., 0][obj])
+      tot_obj += torch.sum(obj)
+      correct_noobj += torch.sum(obj_preds[noobj] == y[i][..., 0][noobj])
+      tot_noobj += torch.sum(noobj)
+
+  return((correct_class/(tot_class_preds+1e-16))*100, 
+         (correct_noobj/(tot_noobj+1e-16))*100, 
+         (correct_obj/(tot_obj+1e-16))*100)
+
+
+def get_datasets(train_loc="/train.csv", test_loc="/test.csv"):
+
+  train_dataset = YOLODataset(
+      config.DATASET + train_loc,
+      transform=config.train_transform,
+      img_dir=config.IMG_DIR,
+      label_dir=config.LABEL_DIR,
+      anchors=config.ANCHORS,
+  )
+
+  test_dataset = YOLODataset(
+      config.DATASET + test_loc,
+      transform=config.test_transform,
+      img_dir=config.IMG_DIR,
+      label_dir=config.LABEL_DIR,
+      anchors=config.ANCHORS,
+      train=False
+  )
+
+  return(train_dataset, test_dataset)
+
+
+
+class YOLOv3Lightning(pl.LightningModule):
+  def __init__(self, dataset=None, lr=config.LEARNING_RATE):
+    super().__init__()
+
+    self.save_hyperparameters()
+
+    self.model = YOLOv3(num_classes=config.NUM_CLASSES)
+    self.lr = lr
+    self.criterion = criterion
+    self.losses = []
+    self.threshold = config.CONF_THRESHOLD
+    self.iou_threshold = config.NMS_IOU_THRESH
+    self.train_idx = 0
+    self.box_format="midpoint"
+    self.dataset = dataset
+    self.criterion = criterion
+    self.accuracy_fn = accuracy_fn
+    self.tot_class_preds, self.correct_class = 0, 0
+    self.tot_noobj, self.correct_noobj = 0, 0
+    self.tot_obj, self.correct_obj = 0, 0
+    self.scaled_anchors = 0
+
+  def forward(self, x):
+    return self.model(x)
+
+  def set_scaled_anchor(self, scaled_anchors):
+      self.scaled_anchors = scaled_anchors
+
+  def on_train_epoch_start(self):
+      # Set a new image size for the dataset at the beginning of each epoch
+      size_idx = random.choice(range(len(config.IMAGE_SIZES)))
+      self.dataset.set_image_size(size_idx)
+      self.set_scaled_anchor((
+          torch.tensor(config.ANCHORS)
+          * torch.tensor(config.S[size_idx]).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)
+      ))
+
+  def on_validation_epoch_start(self):
+      self.set_scaled_anchor((
+          torch.tensor(config.ANCHORS)
+          * torch.tensor(config.S[1]).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)
+      ))
+
+
+  def training_step(self, batch, batch_idx):
+    x, y = batch
+    out = self(x)
+    loss = self.criterion(out, y, self.scaled_anchors)
+
+    self.log('train_loss', loss, prog_bar=True, on_epoch=True, on_step=True, logger=True)
+
+    return loss
+
+  def validation_step(self, val_batch, batch_idx):
+    x, labels = val_batch
+    out = self(x)
+
+    loss = self.criterion(out, labels, self.scaled_anchors)
+    self.log('val_loss', loss, prog_bar=True, on_epoch=True)
+
+    self.evaluate(x, labels, out, 'val')
+
+
+  def evaluate(self, x, y, out, stage=None):
+
+    # Class Accuracy
+    class_accuracy, no_obj_accuracy, obj_accuracy = self.accuracy_fn(y,
+                                                                     out,
+                                                                     self.threshold,
+                                                                     self.correct_class,
+                                                                     self.correct_obj,
+                                                                     self.correct_noobj,
+                                                                     self.tot_class_preds,
+                                                                     self.tot_obj,
+                                                                     self.tot_noobj, )
+    if stage:
+      self.log(f'{stage}_class_accuracy', class_accuracy, prog_bar=True, on_epoch=True, on_step=True, logger=True)
+      self.log(f'{stage}_no_obj_accuracy', no_obj_accuracy, prog_bar=True, on_epoch=True, on_step=True, logger=True)
+      self.log(f'{stage}_obj_accuracy', obj_accuracy, prog_bar=True, on_epoch=True, on_step=True, logger=True)