🎨 [Update] forward proccess, only save usable data
Browse files- yolo/model/yolo.py +4 -3
yolo/model/yolo.py
CHANGED
@@ -58,15 +58,16 @@ class YOLO(nn.Module):
|
|
58 |
layer_idx += 1
|
59 |
|
60 |
def forward(self, x):
|
61 |
-
y =
|
62 |
output = []
|
63 |
-
for layer in self.model:
|
64 |
if isinstance(layer.source, list):
|
65 |
model_input = [y[idx] for idx in layer.source]
|
66 |
else:
|
67 |
model_input = y[layer.source]
|
68 |
x = layer(model_input)
|
69 |
-
|
|
|
70 |
if layer.output:
|
71 |
output.append(x)
|
72 |
return output
|
|
|
58 |
layer_idx += 1
|
59 |
|
60 |
def forward(self, x):
|
61 |
+
y = {0: x}
|
62 |
output = []
|
63 |
+
for index, layer in enumerate(self.model, start=1):
|
64 |
if isinstance(layer.source, list):
|
65 |
model_input = [y[idx] for idx in layer.source]
|
66 |
else:
|
67 |
model_input = y[layer.source]
|
68 |
x = layer(model_input)
|
69 |
+
if hasattr(layer, "save"):
|
70 |
+
y[index] = x
|
71 |
if layer.output:
|
72 |
output.append(x)
|
73 |
return output
|