| """ | |
| Main file for training Yolo model on Pascal VOC and COCO dataset | |
| """ | |
| import torch | |
| from pytorch_lightning import LightningModule, Trainer, seed_everything | |
| import torch | |
| import warnings | |
| warnings.filterwarnings("ignore") | |
| torch.backends.cudnn.benchmark = True | |
| class YOLOTraining(LightningModule): | |
| def __init__(self,model): | |
| super().__init__() | |
| self.model = model | |
| def forward(self, x): | |
| return self.model(x) | |
| if __name__ == "__main__": | |
| num_classes = 20 |