minor changes

YoloV3Lightning.py → yolov3.py

@@ -98,7 +98,7 @@ class ScalePrediction(nn.Module):
         )
 
 
-class YOLOv3LightningModel(pl.LightningModule):
+class YoloV3(pl.LightningModule):
     def __init__(self, in_channels=3, num_classes=20):
         super().__init__()
         self.num_classes = num_classes
@@ -159,13 +159,4 @@ class YOLOv3LightningModel(pl.LightningModule):
                     layers.append(nn.Upsample(scale_factor=2),)
                     in_channels = in_channels * 3
 
-        return layers
-
-
-def sanity_check(model):
-    x = torch.randn((2, 3, cfg.IMAGE_SIZE, cfg.IMAGE_SIZE))
-    out = model(x)
-    assert model(x)[0].shape == (2, 3, cfg.IMAGE_SIZE // 32, cfg.IMAGE_SIZE // 32, cfg.NUM_CLASSES + 5)
-    assert model(x)[1].shape == (2, 3, cfg.IMAGE_SIZE // 16, cfg.IMAGE_SIZE // 16, cfg.NUM_CLASSES + 5)
-    assert model(x)[2].shape == (2, 3, cfg.IMAGE_SIZE // 8, cfg.IMAGE_SIZE // 8, cfg.NUM_CLASSES + 5)
-    print("Success!")
+        return layers