Update model.py

model.py

@@ -163,6 +163,147 @@ class YOLOv3(LightningModule):
 
         return layers
 
+class YoloVersion3(LightningModule):
+    def __init__(self):
+        super(YoloVersion3, self).__init__( )
+        self.save_hyperparameters()
+       # Set our init args as class attributes
+        self.learning_rate=config.LEARNING_RATE
+        #self.config=config
+
+        self.num_classes=config.NUM_CLASSES
+        self.train_csv=config.DATASET + "/train.csv"
+        self.test_csv=config.DATASET + "/test.csv"
+
+        self.loss_fn= YoloLoss()
+        self.scaler = amp.GradScaler()
+        #self.train_transform_function= config.train_transforms
+        #self.in_channels = 3
+        self.model= YOLOv3(num_classes=config.NUM_CLASSES).to(config.DEVICE)
+        self.scaled_anchors = (
+                torch.tensor(config.ANCHORS) * torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)).to(config.DEVICE)
+        #self.register_buffer("scaled_anchors", self.scaled_anchors)
+        self.training_step_outputs = []
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y0, y1, y2 = (
+            y[0],
+            y[1],
+            y[2],
+        )
+        out = self(x)
+        loss = (
+                self.loss_fn(out[0], y0, self.scaled_anchors[0])
+                + self.loss_fn(out[1], y1, self.scaled_anchors[1])
+                + self.loss_fn(out[2], y2, self.scaled_anchors[2])
+        )
+        self.log("train_loss", loss, on_epoch=True, prog_bar=True, logger=True)  # Logging the training loss for visualization
+        self.training_step_outputs.append(loss)
+        return loss
+
+    def on_train_epoch_end(self):
+
+        print(f"\nCurrently epoch {self.current_epoch}")
+        train_epoch_average = torch.stack(self.training_step_outputs).mean()
+        self.training_step_outputs.clear()
+        print(f"Train loss {train_epoch_average}")
+        print("On Train Eval loader:")
+        print("On Train loader:")
+        class_accuracy, no_obj_accuracy, obj_accuracy = check_class_accuracy(self.model, self.train_loader, threshold=config.CONF_THRESHOLD)
+        self.log("class_accuracy", class_accuracy, on_epoch=True, prog_bar=True, logger=True)
+        self.log("no_obj_accuracy", no_obj_accuracy, on_epoch=True, prog_bar=True, logger=True)
+        self.log("obj_accuracy", obj_accuracy, on_epoch=True, prog_bar=True, logger=True)
+
+        if (self.current_epoch>0) and ((self.current_epoch+1) % 6 == 0): # for every 10 epochs we are plotting
+            plot_couple_examples(self.model, self.test_loader, 0.6, 0.5, self.scaled_anchors)
+
+        if (self.current_epoch>0) and (self.current_epoch+1 == self.trainer.max_epochs ): #map calculation across last epoch
+            check_class_accuracy(self.model, self.test_loader, threshold=config.CONF_THRESHOLD)
+            pred_boxes, true_boxes = get_evaluation_bboxes(
+                self.test_loader,
+                self.model,
+                iou_threshold=config.NMS_IOU_THRESH,
+                anchors=config.ANCHORS,
+                threshold=config.CONF_THRESHOLD,
+            )
+            mapval = mean_average_precision(
+                pred_boxes,
+                true_boxes,
+                iou_threshold=config.MAP_IOU_THRESH,
+                box_format="midpoint",
+                num_classes=config.NUM_CLASSES,
+            )
+            print(f"MAP: {mapval.item()}")
+
+            self.log("MAP", mapval.item(), on_epoch=True, prog_bar=True, logger=True)
+
+
+
+    def configure_optimizers(self):
+        optimizer = optim.Adam(
+            self.parameters(),
+            lr=config.LEARNING_RATE,
+            weight_decay=config.WEIGHT_DECAY,
+        )
+
+        self.trainer.fit_loop.setup_data()
+        dataloader = self.trainer.train_dataloader
+
+        EPOCHS = config.NUM_EPOCHS # 40 % of number of epochs
+        lr_scheduler = OneCycleLR(
+            optimizer,
+            max_lr=1E-3,
+            steps_per_epoch=len(dataloader),
+            epochs=EPOCHS,
+            pct_start=5/EPOCHS,
+            div_factor=100,
+            three_phase=False,
+            final_div_factor=100,
+            anneal_strategy='linear'
+        )
+
+        scheduler = {"scheduler": lr_scheduler, "interval" : "step"}
+
+        return [optimizer]
+
+    def setup(self, stage=None):
+        self.train_loader, self.test_loader, self.train_eval_loader = get_loaders(
+            train_csv_path=self.train_csv,
+            test_csv_path=self.test_csv,
+        )
+
+    def train_dataloader(self):
+        return self.train_loader
+
+    def val_dataloader(self):
+        return self.train_eval_loader
+
+    def test_dataloader(self):
+        return self.test_loader
+# if __name__ == "__main__":
+
+#     model = YoloVersion3()
+
+#     checkpoint = ModelCheckpoint(filename='last_epoch', save_last=True)
+#     lr_rate_monitor = LearningRateMonitor(logging_interval="epoch")
+#     trainer = pl.Trainer(
+#                   max_epochs=config.NUM_EPOCHS,
+#                   deterministic=True,
+#                   logger=True,
+#                   default_root_dir="/content/drive/MyDrive/sunandini/Checkpoint/",
+#                   callbacks=[lr_rate_monitor],
+#                   enable_model_summary=False,
+#                   log_every_n_steps=1,
+#                   precision="16-mixed"
+#               )
+#     print("---- Training Started ---- Sunandini ----")
+#     trainer.fit(model)
+#     torch.save(model.state_dict(), 'YOLOv3.pth')
+
 
 if __name__ == "__main__":
     num_classes = 20