Spaces:
Sleeping
Sleeping
File size: 5,318 Bytes
4152812 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
from torchvision import datasets, transforms
import albumentations as Al
from albumentations.pytorch import ToTensorV2
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from torch.optim.lr_scheduler import OneCycleLR
from pytorch_lightning import LightningModule, Trainer, seed_everything
from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
from pytorch_lightning.callbacks.progress import TQDMProgressBar
from pytorch_lightning.loggers import CSVLogger,TensorBoardLogger
from tqdm import tqdm
import torch
import torch.optim as optim
import matplotlib
import cv2
# my files
import utils
import config
from model import YOLOv3
from utils import (
mean_average_precision,
cells_to_bboxes,
get_evaluation_bboxes,
save_checkpoint,
load_checkpoint,
check_class_accuracy,
plot_couple_examples,
accuracy_fn,
get_loaders
)
from loss import YoloLoss
# custom functions for yolo
# loss function for yolov3
loss_fn = YoloLoss()
def model_criterion(out, y,anchors):
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
# accuracy function for yolov3
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)
# pytorch lightning
class LitYolo(LightningModule):
def __init__(self, num_classes=config.NUM_CLASSES, lr=1E-3,weight_decay=config.WEIGHT_DECAY,threshold=config.CONF_THRESHOLD):
super().__init__()
self.save_hyperparameters()
self.model = YOLOv3(num_classes=self.hparams.num_classes)
self.criterion = model_criterion
self.accuracy_fn = accuracy_fn
self.scaled_anchors = (torch.tensor(config.ANCHORS) * torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2))
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
def forward(self, x):
out = self.model(x)
return out
def training_step(self, batch, batch_idx):
x, y = batch
out = self(x)
loss = self.criterion(out,y,self.scaled_anchors)
acc = self.accuracy_fn(y,out,self.hparams.threshold,self.correct_class,
self.correct_obj,
self.correct_noobj,
self.tot_class_preds,
self.tot_obj,
self.tot_noobj)
self.log('train_loss', loss, prog_bar=True, on_step=False, on_epoch=True)
self.log_dict({"class_accuracy": acc[0], "no_object_accuracy": acc[1], "object_accuracy":acc[2]},prog_bar=True,on_step=False, on_epoch=True)
return loss
def evaluate(self, batch, stage=None):
x, y = batch
out = self(x)
loss = self.criterion(out,y,self.scaled_anchors)
acc = self.accuracy_fn(y,out,self.hparams.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}_loss", loss, prog_bar=True)
self.log_dict({"class_accuracy": acc[0], "no_object_accuracy": acc[1], "object_accuracy":acc[2]},prog_bar=True)
def test_step(self, batch, batch_idx):
self.evaluate(batch, "test")
def validation_step(self, batch, batch_idx):
self.evaluate(batch, "val")
def configure_optimizers(self):
optimizer = optim.Adam(self.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay)
scheduler = OneCycleLR(
optimizer,
max_lr= 1E-3,
pct_start = 5/self.trainer.max_epochs,
epochs=self.trainer.max_epochs,
steps_per_epoch=len(train_loader),
div_factor=100,verbose=True,
three_phase=False
)
return ([optimizer],[scheduler]) |