Commit
β’
3650f4f
1
Parent(s):
922be09
minor changes
Browse files
YoloV3Lightning.py β yolov3.py
RENAMED
@@ -98,7 +98,7 @@ class ScalePrediction(nn.Module):
|
|
98 |
)
|
99 |
|
100 |
|
101 |
-
class
|
102 |
def __init__(self, in_channels=3, num_classes=20):
|
103 |
super().__init__()
|
104 |
self.num_classes = num_classes
|
@@ -159,13 +159,4 @@ class YOLOv3LightningModel(pl.LightningModule):
|
|
159 |
layers.append(nn.Upsample(scale_factor=2),)
|
160 |
in_channels = in_channels * 3
|
161 |
|
162 |
-
return layers
|
163 |
-
|
164 |
-
|
165 |
-
def sanity_check(model):
|
166 |
-
x = torch.randn((2, 3, cfg.IMAGE_SIZE, cfg.IMAGE_SIZE))
|
167 |
-
out = model(x)
|
168 |
-
assert model(x)[0].shape == (2, 3, cfg.IMAGE_SIZE // 32, cfg.IMAGE_SIZE // 32, cfg.NUM_CLASSES + 5)
|
169 |
-
assert model(x)[1].shape == (2, 3, cfg.IMAGE_SIZE // 16, cfg.IMAGE_SIZE // 16, cfg.NUM_CLASSES + 5)
|
170 |
-
assert model(x)[2].shape == (2, 3, cfg.IMAGE_SIZE // 8, cfg.IMAGE_SIZE // 8, cfg.NUM_CLASSES + 5)
|
171 |
-
print("Success!")
|
|
|
98 |
)
|
99 |
|
100 |
|
101 |
+
class YoloV3(pl.LightningModule):
|
102 |
def __init__(self, in_channels=3, num_classes=20):
|
103 |
super().__init__()
|
104 |
self.num_classes = num_classes
|
|
|
159 |
layers.append(nn.Upsample(scale_factor=2),)
|
160 |
in_channels = in_channels * 3
|
161 |
|
162 |
+
return layers
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|