RashiAgarwal/YOLO_v3_PyTorchLightning

1"""2Main file for training Yolo model on Pascal VOC and COCO dataset3"""4 5import torch6from pytorch_lightning import LightningModule, Trainer, seed_everything7import torch 8 9import warnings10warnings.filterwarnings("ignore")11 12torch.backends.cudnn.benchmark = True13 14class YOLOTraining(LightningModule):15    def __init__(self,model):16        super().__init__()17        self.model = model18        19    def forward(self, x):20        return self.model(x)21 22if __name__ == "__main__":23    num_classes = 20